Address review comments: remove redundant checks and c_axis property, eliminate code duplication

Co-authored-by: ricardoV94 <[email protected]>

Author: Copilot

pytensor/link/numba/dispatch/extra_ops.py

@@ -37,10 +37,7 @@ def numba_funcify_CumOp(op: CumOp, node: Apply, **kwargs):
     mode = op.mode
     ndim = cast(TensorVariable, node.outputs[0]).ndim
 
-    if axis < 0:
-        axis = ndim + axis
-    if axis < 0 or axis >= ndim:
-        raise ValueError(f"Invalid axis {axis} for array with ndim {ndim}")
+
 
     reaxis_first = (axis, *(i for i in range(ndim) if i != axis))
     reaxis_first_inv = tuple(np.argsort(reaxis_first))

pytensor/tensor/extra_ops.py

@@ -294,7 +294,7 @@ class CumOp(COp):
     __props__ = ("axis", "mode")
     check_input = False
     params_type = ParamsType(
-        c_axis=int_t, mode=EnumList(("MODE_ADD", "add"), ("MODE_MUL", "mul"))
+        axis=int_t, mode=EnumList(("MODE_ADD", "add"), ("MODE_MUL", "mul"))
     )
 
     def __init__(self, axis: int, mode="add"):
@@ -307,9 +307,6 @@ def __init__(self, axis: int, mode="add"):
         self.axis = axis
         self.mode = mode
 
-    @property
-    def c_axis(self) -> int:
-        return self.axis
 
     def make_node(self, x):
         x = ptb.as_tensor_variable(x)
@@ -360,7 +357,7 @@ def c_code(self, node, name, inames, onames, sub):
         fail = sub["fail"]
         params = sub["params"]
 
-        axis_code = f"int axis = {params}->c_axis;\n"
+        axis_code = f"int axis = {params}->axis;\n"
 
         code = (
             axis_code
@@ -432,12 +429,13 @@ def cumsum(x, axis=None):
     .. versionadded:: 0.7
 
     """
+    x = ptb.as_tensor_variable(x)
     if axis is None:
-        return CumOp(axis=0, mode="add")(ptb.as_tensor_variable(x).ravel())
+        x = x.ravel()
+        axis = 0
     else:
-        x = ptb.as_tensor_variable(x)
         axis = normalize_axis_index(axis, x.ndim)
-        return CumOp(axis=axis, mode="add")(x)
+    return CumOp(axis=axis, mode="add")(x)
 
 
 def cumprod(x, axis=None):
@@ -457,12 +455,13 @@ def cumprod(x, axis=None):
     .. versionadded:: 0.7
 
     """
+    x = ptb.as_tensor_variable(x)
     if axis is None:
-        return CumOp(axis=0, mode="mul")(ptb.as_tensor_variable(x).ravel())
+        x = x.ravel()
+        axis = 0
     else:
-        x = ptb.as_tensor_variable(x)
         axis = normalize_axis_index(axis, x.ndim)
-        return CumOp(axis=axis, mode="mul")(x)
+    return CumOp(axis=axis, mode="mul")(x)
 
 
 @_vectorize_node.register(CumOp)

tests/tensor/test_extra_ops.py

@@ -225,7 +225,7 @@ def test_infer_shape(self):
         x = tensor3("x")
         a = np.random.random((3, 5, 2)).astype(config.floatX)
 
-        # Test axis=None using cumsum function (which now handles it symbolically)
+        # Test default axis=None
         self._compile_and_check([x], [cumsum(x)], [a], self.op_class)
 
         for axis in range(-len(a.shape), len(a.shape)):
@@ -234,7 +234,7 @@ def test_infer_shape(self):
     def test_grad(self):
         a = np.random.random((3, 5, 2)).astype(config.floatX)
 
-        # Test axis=None using cumsum/cumprod functions (which now handle it symbolically)
+        # Test default axis=None using cumsum/cumprod functions
         utt.verify_grad(lambda x: cumsum(x), [a])  # Test axis=None for cumsum
         utt.verify_grad(lambda x: cumprod(x), [a])  # Test axis=None for cumprod