Add gallery notebook on vector jacobian products

[ricardoV94]

When we have more powerful pytensor rewrites I would like to make a companion piece showing pytensor deriving the optimized Lop from the naive formulation (maybe with the exception of the transpose)

Incidentally, this PR implements vectorize for CumSum, so it shows up nicely in the jacobian :)

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[Copilot]

Pull Request Overview

This PR adds vectorization support for the cumulative sum operation in PyTensor and includes a new tutorial notebook demonstrating optimized vector-Jacobian products.

• Implement _vectorize_node rule for CumOp to enable batched inputs.
• Add a comprehensive vector_jacobian_product.ipynb showing VJP optimizations across several tensor operations.

Reviewed Changes

Copilot reviewed 2 out of 3 changed files in this pull request and generated no comments.

File	Description
pytensor/tensor/extra_ops.py	Register and implement vectorize_cum_op for CumOp
doc/gallery/gradient/vector_jacobian_product.ipynb	New tutorial notebook illustrating vector-Jacobian products

Comments suppressed due to low confidence (4)

pytensor/tensor/extra_ops.py:477

• [nitpick] There aren’t any unit tests covering the new vectorize_cum_op rule. Consider adding tests for both the axis=None and explicit axis cases to ensure correct batched behavior.

@_vectorize_node.register(CumOp)

pytensor/tensor/extra_ops.py:486

• The function normalize_axis_index is used here but not imported, which will cause a NameError. Please add an import, for example:

from pytensor.tensor.utils import normalize_axis_index

        axis = normalize_axis_index(op.axis, original_x.ndim)

pytensor/tensor/extra_ops.py:490

• Inside the list comprehension you iterate over batch_x but still call batch_x.flatten(...). You likely intended to call x.flatten(...) for each element or otherwise flatten the entire batch tensor directly. Consider:

batch_x_raveled = [x.flatten(ndim=batch_ndim + 1) for x in batch_x]

        batch_x_raveled = [batch_x.flatten(ndim=batch_ndim + 1) for x in batch_x]

doc/gallery/gradient/vector_jacobian_product.ipynb:62

• Spelling mistake: 'Elemtwise' should be 'Elementwise'.

    "## Elemtwise operations"

[codecov]

Codecov Report

❌ Patch coverage is 68.42105% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.53%. Comparing base (12213d0) to head (d049dad).
⚠️ Report is 142 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/tensor/extra_ops.py	60.00%	3 Missing and 3 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1544      +/-   ##
==========================================
+ Coverage   81.49%   81.53%   +0.04%     
==========================================
  Files         232      230       -2     
  Lines       53122    53066      -56     
  Branches     9444     9423      -21     
==========================================
- Hits        43292    43269      -23     
+ Misses       7382     7364      -18     
+ Partials     2448     2433      -15     

Files with missing lines	Coverage Δ
pytensor/tensor/subtensor.py	90.01% <100.00%> (+0.03%)	⬆️
pytensor/tensor/extra_ops.py	88.56% <60.00%> (-0.33%)	⬇️

... and 14 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[ricardoV94]

Link to the docs preview: https://pytensor--1544.org.readthedocs.build/en/1544/gallery/autodiff/vector_jacobian_product.html

And yes the auto snapshoot seems to be erasing the custom pics

[jessegrabowski]

typo: get

Reply via ReviewNB

[jessegrabowski]

looks great!

Did you want to adjust the generate_gallery script to use custom images? Or open an issue for later?

[ricardoV94]

Opened issue #1554