Use vectorized jacobian in minimize ops
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/
Check out this pull request on ![]()
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
Is this PR superseded by https://github.com/pymc-devs/pytensor/pull/1582 ?
I wanted to rerun the notebook example, but that's blocked by the mentioned bug
It has some other changes we may want to cherry-pick to other PRs