Constant fold branches of variadic add/mul

[ricardoV94]

Refactoring and renaming:

• Renamed the local_add_mul_fusion function to flatten_nested_add_mul to more precisely reflect how it works (one could also fuse non-nested add/mul, like the FusionOptimizer does). The function now explicitly tracks add and mul operations instead of relying on generic Elemwise checks. [1] [2] [3]

New optimization for constant folding:

• Introduced a new rewrite function, constant_fold_branches_of_add_mul, which folds constants in add/mul operations when it does not result in higher intermediate memory usage. This optimization is registered in a new sequence database, add_mul_flat_seqopt, which runs before generic elementwise fusion.

The two rewrites are pulled out to a separate database so it's included in JAX rewrites (JAX does not include fusion rewrites). We've found this could help avoding XLA constant fold (CC @lucianopaz)

---

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

[lucianopaz]

I don’t know why the test fails but it looks like the fusion rewrite is applied only once. Maybe the equilibrium rewrite that you took out should be added back in?

[lucianopaz]

@ricardoV94, I just went through your branch's code and found that the error is coming from the fact that the TestFusion class is including: "canonicalize", "fusion", and "inplace" rewrite databases. The add_mul flatten and fusion rewrites that you moved or added here are only included in "fast_run". My question then is whether your rewrites should also be added to fusion, or if you only want to add the fast_run database rewrites to the TestFusion includes?

[ricardoV94]

This was the change needed to get the test to pass @lucianopaz

[ricardoV94]

@lucianopaz I came to the same conclusion, I just added the rewrite explicitly. Mentioned in an inline comment above

[lucianopaz]

LGTM, thanks @ricardoV94 !

[codecov]

Codecov Report

❌ Patch coverage is 95.34884% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.13%. Comparing base (5a462e9) to head (70db72e).
⚠️ Report is 136 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/tensor/rewriting/elemwise.py	95.34%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@ Coverage Diff @@ ## main #1422 +/- ## ========================================== + Coverage 82.11% 82.13% +0.01%  ========================================== Files 211 211 Lines 49743 49773 +30 Branches 8824 8830 +6 ========================================== + Hits 40847 40879 +32  + Misses 6715 6714 -1  + Partials 2181 2180 -1 

Files with missing lines	Coverage Δ
pytensor/tensor/rewriting/elemwise.py	91.77% <95.34%> (+0.74%)	⬆️

🚀 New features to boost your workflow:

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