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