diff --git a/pytensor/link/numba/dispatch/elemwise.py b/pytensor/link/numba/dispatch/elemwise.py index e6cbab4c4c..df8aeefbdd 100644 --- a/pytensor/link/numba/dispatch/elemwise.py +++ b/pytensor/link/numba/dispatch/elemwise.py @@ -841,7 +841,7 @@ def dimshuffle_inner(x, shuffle): @numba_basic.numba_njit def dimshuffle_inner(x, shuffle): - return np.reshape(x, ()) + return np.reshape(np.ascontiguousarray(x), ()) # Without the following wrapper function we would see this error: # E No implementation of function Function() found for signature: diff --git a/tests/link/numba/test_elemwise.py b/tests/link/numba/test_elemwise.py index 293538a347..2e5bb5c1e4 100644 --- a/tests/link/numba/test_elemwise.py +++ b/tests/link/numba/test_elemwise.py @@ -218,6 +218,17 @@ def test_Dimshuffle_returns_array(): assert out.ndim == 0 +def test_Dimshuffle_non_contiguous(): + """The numba impl of reshape doesn't work with + non-contiguous arrays, make sure we work around that.""" + x = at.dvector() + idx = at.vector(dtype="int64") + op = pytensor.tensor.elemwise.DimShuffle([True], []) + out = op(at.specify_shape(x[idx][::2], (1,))) + func = pytensor.function([x, idx], out, mode="NUMBA") + assert func(np.zeros(3), np.array([1])).ndim == 0 + + @pytest.mark.parametrize( "careduce_fn, axis, v", [