pymc-devs
WIP: Add rewrite to fuse nested BlockDiag Ops
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
- [x] Closes #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
- [x] New feature / enhancement
- [ ] Bug fix
- [ ] Documentation
- [ ] Maintenance
- [ ] Other (please specify):
📚 Documentation preview 📚: https://pytensor--1671.org.readthedocs.build/en/1671/
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.
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.
@jessegrabowski ! Added a rewrite to fuse nested BlockDiagonal ops and updated test_linalg.py with a test for nested BlockDiagonal fusion.
@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.