Help with scan in logp calculation

Hi all,

I am having some difficulty understanding loops using scan. I am building a logp for an object X = (X_1, …, X_d) with a CustomDist that takes as a parameter an array called samples of shape (N, d). These samples are drawn from some joint probability on (X_i), and I am writing this CustomDist to approximate that joint probability.

As part of the logp calculation, there is a step where each X_i has a contribution f_i(X_i), where f_i is the marginal pdf on the i-th component, which I estimate using a histogram of samples[:, i-1] called marginals_hists with shape (d, 2, num_bins). To do calculate f_i(X_i), I have been trying to use scan since the function f_i is different for each i. I defined a function _get_marginal_cdf which performs the calculation of f_i(X_i) for a single dimension, and then loop over dimensions using scan, as below.

def _marginal_cdf_forward(value, marginals_hists):
    output, updates = pytensor.scan(fn=_get_marginal_cdf,
                                    sequences=[value, marginals_hists])
    transformed_value = pytensor.function(inputs=[value.T, marginals_hists],
                                          outputs=output.T,
                                          updates=updates)
    return transformed_value

However, when I try to call logp during sampling, I get the following error (including the last part of the error trace starting from the call to logp).

Cell In[42], line 68, in MvExtendedInterpolatedWithEmpiricalBetaCopula.logp(value, samples, ranks, marginals_hists)
     67 def logp(value, samples, ranks, marginals_hists):
---> 68     transformed_value = _marginal_cdf_forward(value, marginals_hists.data)
     69     ebc_density = empirical_beta_copula_density(transformed_value, samples.data)
     70     marginals_densities = _marginal_logpdfs(value, marginals_hists.data)

Cell In[41], line 248, in _marginal_cdf_forward(value, marginals_hists)
    247 def _marginal_cdf_forward(value, marginals_hists):
--> 248     output, updates = pytensor.scan(fn=_get_marginal_cdf,
    249                                     sequences=[value, marginals_hists])
    250     transformed_value = pytensor.function(inputs=[value.T, marginals_hists],
    251                                           outputs=output.T,
    252                                           updates=updates)
    253     return transformed_value

File ~/.conda/envs/pymc_env/lib/python3.12/site-packages/pytensor/scan/basic.py:681, in scan(fn, sequences, outputs_info, non_sequences, n_steps, truncate_gradient, go_backwards, mode, name, profile, allow_gc, strict, return_list)
    678 else:
    679     actual_n_steps = pt.as_tensor(n_steps)
--> 681 scan_seqs = [seq[:actual_n_steps] for seq in scan_seqs]
    682 # Conventions :
    683 #   mit_mot = multiple input taps, multiple output taps ( only provided
    684 #             by the gradient function )
   (...)
    688 
    689 # MIT_MOT -- not provided by the user only by the grad function
    690 n_mit_mot = 0

File ~/.conda/envs/pymc_env/lib/python3.12/site-packages/pytensor/tensor/variable.py:37, in _tensor_py_operators.__index__(self)
     36 def __index__(self):
---> 37     raise TypeError(
     38         "TensorVariable cannot be converted to Python integer. "
     39         "Call `.astype(int)` for the symbolic equivalent."
     40     )

TypeError: TensorVariable cannot be converted to Python integer. Call `.astype(int)` for the symbolic equivalent.

Any help addressing this issue would be much appreciated!

Couple of things:

  • What is the type of value, marginal_hists? Are they PyTensor variables?
  • You probably don’t want to create a PyTensor function internally
  • Even if you wanted to create one you can say inputs=[value.T, ...]. Inputs must be in the graph of outputs, and value.T is by definition a new variable.
  • I pass alongvalue from the call to logp so it is a TensorVariable. Originally, marginal_hists is a pt.constant, and then I pass marginal_hists.data to _marginal_cdf_forward in logp (see below).
  • Originally, I tried to just do a for loop on the dimension of the distribution by taking samples.shape[1], but it was not happy with expressions like value[:, d]. Is there another way to loop? I include at the bottom my code for _get_marginal_cdf which is modeled off the code for pm.Interpolated.
  • Can I ask what you mean by the graph of outputs?

Here is my logp:

def logp(value, samples, ranks, marginals_hists):
    transformed_value = _marginal_cdf_forward(value, marginals_hists.data)
    ebc_density = empirical_beta_copula_density(transformed_value, samples.data)
    marginals_densities = _marginal_logpdfs(value, marginals_hists.data)
        
    return pt.log(ebc_density) + pt.math.sum(_marginal_logpdfs, axis=1)
def _get_marginal_cdf(value_column, marginals_hists_dim):
    interp = InterpolatedUnivariateSpline(marginals_hists_dim[0], marginals_hists_dim[1], k=1, ext="zeros")
    Z = interp.integral(marginals_hists_dim[0, ..., 0], marginals_hists_dim[0, ..., -1])
    interp_cdf = interp.antiderivative()
    
    interp_cdf_op = SplineWrapper(interp_cdf)
    Z = pt.constant(Z)
    
    return pt.log(interp_cdf_op(value_column) / Z)

I’ve been playing around with how to pass marginals_hists since there seems to be a difficulty somewhere between using a TensorConstant marginals_hists when I need to take a slice (which creates a TensorVariable) and getting to a 1d numpy array to feed into InterpolatedUnivariateSpline. I have tried two cases so far:

  1. I get the ndarray for marginals_hists first, but run into an issue since index is a Tensor object, or
  2. I have marginals_hists[index, 0] as a TensorVariable, but have some difficulty getting the ndarray to pass to scipy.

Do I just use eval() ? Alternatively, is there a way to keep marginals_hists as an ndarray through all of this and deal somehow with index being some Tensor object?

So the problem is InterpolatedUnivariateSplite requires numpy array (doesn’t work with TensorVariables)?

Alternatively, is there a way to keep marginals_hists as an ndarray through all of this and deal somehow with index being some Tensor object?

For the specific question of indexing with a tensor on a numpy array you can do pt.as_tensor(array)[tensor_index].

Regarding .eval(), that is only safe if you’re really dealing with constants, if it’s model parameters or expressions that combine model parameters and constants you shouldn’t use eval or you’ll have a wrong model in the end.

I have to look at your example more careful to give a proper holistic reply. Treat my comments as the myopic responses they are.