pytensor icon indicating copy to clipboard operation
pytensor copied to clipboard

Use vectorized jacobian in minimize ops

Open ricardoV94 opened this issue 4 months ago • 5 comments

This PR uses the more new experimental vectorized jacobian in the minimize/root Ops, with the goal of getting it more tested, as it won't be breaking any people's code.

The existing notebook revealed we were missing some rewrites to avoid Blockwise(IncSubtensor), that this PR also addresses. In the printed example we can see that the vectorized jacobian replaces a Scan with two IncSubtensors by two IncSubtensors on an expanded buffer. This should be much more performant for this sort of small graphs.

Either my chages introduce or reveal a bug on the notebook, so this is still a draft


📚 Documentation preview 📚: https://pytensor--1558.org.readthedocs.build/en/1558/

ricardoV94 avatar Jul 28 '25 14:07 ricardoV94

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

This is hitting https://github.com/pymc-devs/pytensor/issues/1550 I don't know if because I introduced a bug that triggers it or what not. Let's fix that first and then come back to this

ricardoV94 avatar Jul 28 '25 15:07 ricardoV94

Is this PR superseded by https://github.com/pymc-devs/pytensor/pull/1582 ?

jessegrabowski avatar Oct 16 '25 14:10 jessegrabowski

I wanted to rerun the notebook example, but that's blocked by the mentioned bug

ricardoV94 avatar Oct 16 '25 16:10 ricardoV94

It has some other changes we may want to cherry-pick to other PRs

ricardoV94 avatar Oct 16 '25 16:10 ricardoV94