Reshaping within PYMC4

I’m implementing a GP model where I want an interaction between the kernel for space and time. I’m using LatentKron for this, e.g.:

gp = pm.gp.LatentKron( cov_funcs=[cov_func_time, cov_func_space] )

The result of this gives a flattened output of the first 10 timepoints for the first location, followed by the second 10 timepoints for the second location, etc.

I want to reshape this (PYMC4) so that it gives an output shape where each column is the spatial location, and each row is time. How do I go about doing that in PYMC please?

The purpose of this is that I can index the spatial location in a hierarchical model.

Or, if anyone has suggestions of how to use the output of LatentKron in a hierarchical model, that would be great.

You can rehape pymc variables by calling the method .reshape just like Numpy. For example

x = pm.Normal("x", shape=10)
x = x.reshape((2, 5))

Thanks!