Using pytensor random number generator with scan for Simulator class

Questions v5
2

So there is a major enough update that I decided to add this as a new message. About the error that occurs at

File ~/miniconda3/envs/ACORG/lib/python3.10/site-packages/pytensor/tensor/random/op.py:401, in RandomVariable.perform(self, node, inputs, outputs)
    398     rng = copy(rng)
    400 outputs[0][0] = rng
--> 401 outputs[1][0] = np.asarray(
    402     self.rng_fn(rng, *args, None if size is None else tuple(size)),
    403     dtype=self.dtype,
    404 )

I have identified that this first of all occurs for the simulator RV. There are actually two problems that occur here. One is shape related so won’t dwell on that (also not cause of the error I mentioned above). The other one is I assumed, like in normal sampling, that everything was done tensorially and the returned result should be tensors. But seems like it is not the case for the simulator class. In fact the dtype in the code above is float where as self.rng returns tensor hence the “setting an array element with a sequence” error. I fixed the shape and dtype issue by changing_WF_sim as follows (basically adding an eval in the end although it seems I could have done everything in numpy to):

def _WFM_sim(rng, f, m, initial_state, nsims, size=None):

  nvariants = initial_state.size
  N = initial_state.sum()
  rng = pt.shared(rng)
  initial_state = initial_state[0]


  def transition(*args):

    state,  f, m = args

    f = f.flatten()
    m = m.flatten()[0]

    p = _calculate_P_pyt(initial_state, f, m, nvariants)
    next_state = pm.Multinomial.dist(N, p, rng=rng)

    return next_state

  result, updates = pt.scan(transition,
                            outputs_info=initial_state,
                            non_sequences=[f, m],
                            n_steps=nsims)

  return result.eval()

Now the code does not give any errors WHEN evaluating the simulator rv though it suspiciously returns an array with identical states so will need to think on that a bit):

array([[  11.,   39., 9860.,  278.,   22.],
       [  11.,   39., 9860.,  278.,   22.],
       [  11.,   39., 9860.,  278.,   22.],
       [  11.,   39., 9860.,  278.,   22.],
       [  11.,   39., 9860.,  278.,   22.],
       [  11.,   39., 9860.,  278.,   22.],
       [  11.,   39., 9860.,  278.,   22.],
       [  11.,   39., 9860.,  278.,   22.],
       [  11.,   39., 9860.,  278.,   22.]])

However the simulator still does not run to completion it also gives the following error (I can give the full stack if that would be helpful, it happens inside run_chains):

ValueError: Input dimension mismatch: (input[%i].shape[%i] = %lld, input[%i].shape[%i] = %lld)
Apply node that caused the error: Composite{sqr((i0 - i1))}(s{[[0.0000e+ ... 4000e+01]]}, Simulator_s_rv{"(),(),(),()->()"}.out)
Toposort index: 8
Inputs types: [TensorType(float64, shape=(10, 5)), TensorType(float64, shape=(10, 5))]
Inputs shapes: [(10, 5), (9, 5)]
Inputs strides: [(40, 8), (40, 8)]
Inputs values: ['not shown', 'not shown']
Inputs type_num: [12, 12]
Outputs clients: [[Sum{axes=None}(Composite{sqr((i0 - i1))}.0)]]

This one is a bit more harder to diagnose because the error stack just seems to happen inside run_chains functions. But the error suggests that it happens when computing the error between data and observed since Composite{sqr((i0 - i1))}. I will see if I can fix it. Changing sum_stat to “mean” seems to make the chains run forever without any progress, I suppose that is because of the slow eval. I will rewrite in numpy or compile it tomorrow to see if it improves anything.

Makes me wonder if I should just drop the simulator and write everything with potentials as in