diff --git a/pytensor/link/numba/dispatch/extra_ops.py b/pytensor/link/numba/dispatch/extra_ops.py index b59c2ac08f..d616277eb8 100644 --- a/pytensor/link/numba/dispatch/extra_ops.py +++ b/pytensor/link/numba/dispatch/extra_ops.py @@ -49,7 +49,7 @@ def numba_funcify_CumOp(op: CumOp, node: Apply, **kwargs): @numba_basic.numba_njit def cumop(x): - return np.cumsum(x, axis=axis) + return np.cumsum(x) else: @@ -73,7 +73,7 @@ def cumop(x): @numba_basic.numba_njit def cumop(x): - return np.cumprod(x, axis=axis) + return np.cumprod(x) else: diff --git a/tests/link/numba/test_extra_ops.py b/tests/link/numba/test_extra_ops.py index 63884233fd..abc0bf3f59 100644 --- a/tests/link/numba/test_extra_ops.py +++ b/tests/link/numba/test_extra_ops.py @@ -58,6 +58,17 @@ def test_Bartlett(val): 1, "mul", ), + # Regression tests for https://github.com/pymc-devs/pytensor/issues/1689 + ( + (pt.vector(), np.arange(6, dtype=config.floatX)), + 0, + "add", + ), + ( + (pt.vector(), np.arange(6, dtype=config.floatX)), + 0, + "mul", + ), ], ) def test_CumOp(val, axis, mode):