Open fonnesbeck opened 1 year ago
Any change you can give us a MWE?
I think I found it (got the model from Chris on slack):
with pm.Model() as model:
x = pm.Normal("x", shape=(3, 5, 7))
y = at.cumsum(x, axis=-1)
pm.Deterministic("y", y[np.array([2]), np.array([4]), np.array([6])])
func = aesara.function([model.x], model.y, mode="NUMBA")
func(np.zeros((3, 5, 7)))
Seems to be an issue in the numba impl of cumsum
Update:
Directly using aesara
x = at.dtensor3("x")
y = at.cumsum(x, axis=-1)
func = aesara.function([x], y, mode="NUMBA")
func(np.ones((3, 5, 7))).shape
# (5, 7, 3)
Update 2:
I think I figured it out. The problem is that the code assumes that x.transpose(*order).transpose(*order)
recreates the original array order, which is not in general true when rank>2. Working on a fix...
@aseyboldt is this fixed on the PyTensor side?
Running a large model (too slow to be sampled with C backend) results in an
IndexError
when using the Numba backend via:aesara.config.mode = "NUMBA"
The model runs with the C backend, but very slowly. Note that I get the same error with
nutpie
, suggesting that it is an issue with Numba.Versions and main components