Handle `axis=None` symbolically instead of within CumOp #1574

Copilot · 2025-08-09T03:35:54Z

closes #1549

pytensor/tensor/extra_ops.py

tests/tensor/test_extra_ops.py

ricardoV94 · 2025-08-09T04:01:53Z

@copilot address review comments

Copilot · 2025-08-09T04:10:46Z

@copilot address review comments

All review comments have been addressed in commit d3bfb15:

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

pytensor/link/numba/dispatch/extra_ops.py

pytensor/tensor/extra_ops.py

tests/tensor/test_extra_ops.py

ricardoV94 · 2025-08-12T15:37:18Z

@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 · 2025-08-12T15:44:49Z

@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 · 2025-09-30T11:58:27Z

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

jessegrabowski

small nitpicks

jessegrabowski · 2025-10-01T02:03:00Z

pytensor/tensor/extra_ops.py

+        if not isinstance(axis, int):
+            raise TypeError("axis must be an integer.")
+        if axis < 0:
+            raise ValueError("axis must be non-negative.")


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.

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 · 2025-10-01T02:03:45Z

pytensor/tensor/extra_ops.py

-            out_type = vector(dtype=x.dtype)  # Flatten
-        elif self.axis >= x.ndim or self.axis < -x.ndim:
+        if self.axis >= x.ndim:
            raise ValueError(f"axis(={self.axis}) out of bounds")


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

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

Updated the error message as suggested (ce1dfac).

jessegrabowski · 2025-10-01T02:06:43Z

pytensor/tensor/extra_ops.py

 def cumsum(x, axis=None):
    """Return the cumulative sum of the elements along a given `axis`.

    This wraps ``numpy.cumsum``.


Not really?

What not really :)?

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

It does, but it's a detail. Shouldn't be part of the docs

jessegrabowski · 2025-10-17T14:21:04Z

Anything blocking this one?

ricardoV94 · 2025-10-19T20:54:17Z

@copilot can you rebase, resolve the conflicts and address the review comments

Copilot · 2025-10-19T21:10:22Z

@copilot can you rebase, resolve the conflicts and address the review comments

I've completed the rebase onto origin/main, resolved all conflicts, and addressed the review comment about improving the error message formatting (ce1dfac).

The changes are ready locally but I'm unable to force push after the rebase due to tooling limitations. The rebased branch needs a force push to update the remote, which the available tools don't support. You may need to manually push or provide access for force pushing.

codecov · 2025-10-20T13:00:55Z

Codecov Report

❌ Patch coverage is 70.27027% with 11 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.61%. Comparing base (792bd04) to head (b0dbc31).
⚠️ Report is 5 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/tensor/extra_ops.py	75.86%	3 Missing and 4 partials ⚠️
pytensor/link/numba/dispatch/extra_ops.py	66.66%	0 Missing and 2 partials ⚠️
pytensor/link/pytorch/dispatch/extra_ops.py	0.00%	2 Missing ⚠️

❌ Your patch status has failed because the patch coverage (70.27%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1574      +/-   ##
==========================================
- Coverage   81.64%   81.61%   -0.03%     
==========================================
  Files         242      242              
  Lines       53573    53523      -50     
  Branches     9453     9433      -20     
==========================================
- Hits        43737    43681      -56     
- Misses       7358     7364       +6     
  Partials     2478     2478

Files with missing lines	Coverage Δ
pytensor/link/numba/dispatch/extra_ops.py	`92.19% <66.66%> (-3.33%)`	⬇️
pytensor/link/pytorch/dispatch/extra_ops.py	`82.75% <0.00%> (+10.03%)`	⬆️
pytensor/tensor/extra_ops.py	`87.93% <75.86%> (-0.62%)`	⬇️

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Copilot AI assigned Copilot and ricardoV94 Aug 9, 2025

Copilot AI changed the title ~~[WIP] Fix https://github.com/pymc-devs/pytensor/issues/1549~~ Handle special ravelling behavior of CumOp symbolically Aug 9, 2025

Copilot AI requested a review from ricardoV94 August 9, 2025 03:50

ricardoV94 requested changes Aug 9, 2025

View reviewed changes

pytensor/tensor/extra_ops.py Outdated Show resolved Hide resolved

pytensor/tensor/extra_ops.py Outdated Show resolved Hide resolved

tests/tensor/test_extra_ops.py Outdated Show resolved Hide resolved

Copilot AI requested a review from ricardoV94 August 9, 2025 04:11

ricardoV94 reviewed Aug 9, 2025

View reviewed changes

Copilot AI requested a review from ricardoV94 August 9, 2025 04:33

ricardoV94 changed the title ~~Handle special ravelling behavior of CumOp symbolically~~ CumOp no longer accepts axis=None Aug 9, 2025

ricardoV94 changed the title ~~CumOp no longer accepts axis=None~~ Handles axis=None symbolically instead of within CumOp Aug 9, 2025

ricardoV94 added major Op implementation refactor labels Aug 9, 2025

ricardoV94 marked this pull request as ready for review August 9, 2025 14:14

ricardoV94 force-pushed the copilot/fix-4ded62d7-b0ca-436b-9c1f-545c389765f5 branch from aad9f7d to 1cf93c3 Compare August 19, 2025 10:00

ricardoV94 mentioned this pull request Sep 28, 2025

Use incsubtesor in gradient of repeat #1621

Merged

ricardoV94 force-pushed the copilot/fix-4ded62d7-b0ca-436b-9c1f-545c389765f5 branch from 1cf93c3 to a09b463 Compare September 28, 2025 19:56

ricardoV94 requested a review from jessegrabowski September 30, 2025 11:57

ricardoV94 changed the title ~~Handles axis=None symbolically instead of within CumOp~~ Handle axis=None symbolically instead of within CumOp Sep 30, 2025

jessegrabowski approved these changes Oct 1, 2025

View reviewed changes

ricardoV94 force-pushed the copilot/fix-4ded62d7-b0ca-436b-9c1f-545c389765f5 branch from a09b463 to 9b5de87 Compare October 20, 2025 12:13

Implement axis=None raveling behavior symbolically in CumOp

b0dbc31

ricardoV94 force-pushed the copilot/fix-4ded62d7-b0ca-436b-9c1f-545c389765f5 branch from 9b5de87 to b0dbc31 Compare October 20, 2025 12:43

ricardoV94 merged commit 49f76da into main Oct 20, 2025
57 of 58 checks passed

ricardoV94 deleted the copilot/fix-4ded62d7-b0ca-436b-9c1f-545c389765f5 branch October 20, 2025 14:46

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

Handle axis=None symbolically instead of within CumOp #1574

Handle axis=None symbolically instead of within CumOp #1574

Uh oh!

Conversation

Copilot AI commented Aug 9, 2025 • edited by ricardoV94 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ricardoV94 commented Aug 9, 2025

Uh oh!

Copilot AI commented Aug 9, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ricardoV94 commented Aug 12, 2025

Uh oh!

Copilot AI commented Aug 12, 2025

Uh oh!

ricardoV94 commented Sep 30, 2025

Uh oh!

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

jessegrabowski Oct 1, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Oct 1, 2025

Choose a reason for hiding this comment

Uh oh!

jessegrabowski Oct 1, 2025

Choose a reason for hiding this comment

Uh oh!

jessegrabowski Oct 1, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Oct 19, 2025

Choose a reason for hiding this comment

Uh oh!

jessegrabowski Oct 1, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Oct 1, 2025

Choose a reason for hiding this comment

Uh oh!

jessegrabowski Oct 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Oct 19, 2025

Choose a reason for hiding this comment

Uh oh!

jessegrabowski commented Oct 17, 2025

Uh oh!

ricardoV94 commented Oct 19, 2025

Uh oh!

Copilot AI commented Oct 19, 2025

Uh oh!

codecov bot commented Oct 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Handle `axis=None` symbolically instead of within CumOp #1574

Handle `axis=None` symbolically instead of within CumOp #1574

Copilot AI commented Aug 9, 2025 •

edited by ricardoV94

Loading

jessegrabowski Oct 17, 2025 •

edited

Loading

codecov bot commented Oct 20, 2025 •

edited

Loading