Handle axis=None symbolically instead of within CumOp

[Copilot]

Fixes #1549

Summary

This PR refactors the cumulative sum/product operations to handle the special ravelling behavior of axis=None symbolically, making the code cleaner and more maintainable.

Problem

Previously, when cumsum(x, axis=None) or cumprod(x, axis=None) was called on a matrix, the CumOp class internally handled the flattening with special logic scattered throughout:

• Special cases in make_node(), grad(), infer_shape(), and c_code()
• Backend-specific axis=None handling in PyTorch, JAX, and Numba dispatchers
• Complex conditional logic that made the code harder to maintain

Solution

The refactoring separates concerns by handling the ravelling symbolically:

Before:

cumsum(x, axis=None) → CumOp(axis=None)(x)

After:

cumsum(x, axis=None) → CumOp(axis=0)(flatten(x))

Key Changes

1. Modified cumsum/cumprod functions: When axis=None, explicitly flatten the input first, then apply cumsum/cumprod with axis=0

2. Simplified CumOp class:

   • Constructor now requires an integer axis (no more axis=None)
   • Removed all special axis=None handling from make_node, grad, infer_shape, c_code
   • Simplified vectorization logic

3. Updated backend dispatchers: Removed axis=None handling from PyTorch, JAX, and Numba backends since CumOp now always receives integer axes

4. Updated tests: Modified test cases to use the new simplified API while maintaining coverage of axis=None behavior through the public functions

Benefits

• Separation of concerns: Flattening is explicit in the symbolic graph
• Simpler implementation: CumOp only handles specific integer axes
• Better maintainability: No more special cases scattered throughout the codebase
• Clearer symbolic representation: Graph explicitly shows flatten → cumsum chain
• Optimization friendly: Graph rewrites can now optimize flatten and cumsum operations independently

Backward Compatibility

✅ All existing user code continues to work unchanged
✅ Same computational results for all operations
✅ Same gradient behavior
✅ Same output shapes and types

The refactoring only changes the internal implementation - the public API remains identical.

Example

import pytensor.tensor as pt
from pytensor.tensor.extra_ops import cumsum

x = pt.matrix('x')

# This still works exactly the same from user perspective
result = cumsum(x, axis=None)

# But internally now creates: cumsum(flatten(x), axis=0) 
# instead of the old: CumOp(axis=None)(x)

This makes the code more maintainable while preserving all existing functionality.

---

💬 Share your feedback on Copilot coding agent for the chance to win a $200 gift card! Click here to start the survey.

[ricardoV94]

@copilot address review comments

[Copilot]

@copilot address review comments

All review comments have been addressed in commit d3bfb15:

1. Updated cumsum/cumprod to use .ravel() and simplified to one line
2. Added axis normalization in helper functions and simplified negative axis handling
3. Updated test to use self.op_class

[ricardoV94]

@copilot some tests FAILED tests/tensor/test_extra_ops.py::TestCumOp::test_grad are failing because they're trying to use negative axis at the Op level. Remove those cases from the failing tests

Other tests may be failing, check the logs to be sure it's not only this one

[Copilot]

@copilot some tests FAILED tests/tensor/test_extra_ops.py::TestCumOp::test_grad are failing because they're trying to use negative axis at the Op level. Remove those cases from the failing tests

...

Fixed the failing tests by removing negative axis cases from CumOp tests. Updated both TestCumOp.test_grad and the Numba test_CumOp to only use non-negative axis values when directly instantiating the CumOp class, since the constructor now rejects negative axis values (aad9f7d).

[ricardoV94]

@jessegrabowski requested your review because copilot PRs are not really someone else's?

[jessegrabowski]

small nitpicks

[jessegrabowski]

The similar change for repeat told the user to go use the pt.repeat helper in your axis is negative or None, that would be nice here too.

[ricardoV94]

cumop is more annoying because user could be using cumsum or cumprod. Anyway that was a bit overkill in that PR, nobody is instantiating Ops for themselves when using PyTensor

[jessegrabowski]

Suggested change

	raise ValueError(f"axis(={self.axis}) out of bounds")
	raise ValueError(f"axis(={self.axis}) out of bounds for tensor with {x.ndim} dimensions")

[jessegrabowski]

You didn't change this but I saw it, and the robot is doing the work anyway.

[jessegrabowski]

Not really?

[ricardoV94]

What not really :)?

[jessegrabowski]

It doesn't wrap numpy.cumsum, it implements similar functionality though.

[jessegrabowski]

Anything blocking this one?