Add rewrite to fuse nested BlockDiag Ops

[eby0303]

Description

This is a draft PR for issue #1593.
I’m setting up the local environment and exploring how to implement a rewrite that fuses nested BlockDiag Ops into a single one.
I’ll update this PR with code once the setup is complete and I have an initial version of the rewrite.

Related Issue

• Closes Add rewrite to fuse nested BlockDiag Ops #1593

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

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

[eby0303]

Hey @jessegrabowski ! As suggested in the issue discussion, I’ve opened this draft PR to start working on the BlockDiag rewrite.
I’ll be setting up locally. Please feel free to share any tips or guidance for where to begin
Will update this PR as I make progress.

[jessegrabowski]

I'd suggest you work in a test-driven way. Add a test to tests/tensor/rewriting/test_linalg.py with a simple nested blockwise and count the number of BlockDiag ops, and assert that there is only 1. Confim that this test fails. Then add a rewrite to tensor/rewriting/linalg.py that looks for a blockwise with a blockwise inside, and if so merges them.

For an example of how to count ops in a graph for the test, look here (BUT the whole class is overkill for your case, just take the pieces from it and write an inline version).

For a good rewrite template to get you started, I think this one is pretty readable. You will need to 1) check that the input is a BlockDiag op, 2) check that at least one of the inputs to the BlockDiag is a BlockDiag, 3) pull out the inputs from the inner BlockDiag, 4) make a new BlockDiag with n_inputs = old_n_inputs + 1 and return it, passing in all 3 inputs.

[eby0303]

@jessegrabowski ! Added a rewrite to fuse nested BlockDiagonal ops and updated test_linalg.py with a test for nested BlockDiagonal fusion.

[jessegrabowski]

Really great first pass. You're missing some of the boiler plate around rewrites, have a look here (or anywhere in this file really) to see how to register a rewrite, and how to tell it which Op to track (pay attention to the decorators).

You also need to use pytensor.function to compile your block diagonal graph and check that the rewrite was triggered, rather than calling it directly.

[eby0303]

@jessegrabowski Added a rewrite to fuse nested BlockDiagonal ops into a single fused instance and included tests to verify fusion behavior, n_inputs, and output shape.

[ricardoV94]

Oops I rebase-merged instead of squash. Anyway thanks @eby0303 and @jessegrabowski