Ignore axis argument in numba CumOp when input is 1d

Author: jessegrabowski

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:
 

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):