Reshaping a numpy array in a pytensor compiled function

If I use a series of operations like

array = pm.Normal( ... )
array_reshaped = array[:, np.newaxis]

and then go on to use array_reshaped later in the model, is that expected to cause problems for PyMC/PyTensor?

In other words, when the documentation for PyTensor says that it has tight integration with numpy, does that extend to these kinds of operations? Or do I have to use the PyTensor equivalent? And if so, what would it be?

Apologies if this isn’t clear, and thanks in advance for any guidance!

Yes you can do that. The same indexing operations are allowed in PyTensor variables as in numpy arrays.

To debug and check you are getting the expected shapes you can call pm.draw(array_reshaped) to get a random draw.

1 Like

Thank you, @ricardoV94!