pymc-devs / pymc-devs/pytensor
ENH: Extend linear algebra graph rewrites to cover additional inverse and matmul patterns
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Before
Currently, some expressions involving matrix inverse and multiplication may appear in computation graphs, such as:
pt.matmul(pt.linalg.inv(A), B)
pt.matmul(A, pt.linalg.inv(B))
pt.linalg.inv(A).T @ b
While PyTensor already includes rewrites like `inv(A) @ b → solve(A, b)` for certain Dot patterns, similar optimizations may not apply consistently across all equivalent matmul or transpose-based expressions.
As a result, graphs may still contain explicit matrix inverse operations that could otherwise be rewritten into more efficient solve-based formulations.
After
Extend the existing linear algebra rewrite system to detect and optimize additional equivalent patterns involving matrix inverse, transpose, and matrix multiplication.
For example, rewrite patterns such as:
matmul(inv(A), B) → solve(A, B)
matmul(A, inv(B)) → solve(B.T, A.T).T
These rewrites would improve numerical stability, reduce computational cost, and simplify the computation graph.
Context for the issue:
PyTensor already provides several linear algebra graph rewrites in pytensor/tensor/rewriting/linalg.py, including optimizations involving matrix inverse and solve operations.
Extending rewrite coverage to additional equivalent patterns involving matmul, transpose, and inverse would further improve graph optimization and performance for linear algebra workloads.
This enhancement would align with PyTensor’s existing rewrite infrastructure and support more comprehensive optimization of symbolic linear algebra expressions.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in pytensor/tensor/rewriting/linalg.py and inspect the existing inverse and solve rewrites, especially the Dot patterns mentioned in the issue. Trace how equivalent matmul and transpose expressions are represented, then extend coverage for the listed patterns and verify that the resulting graphs use solve-based formulations instead of explicit inverse operations.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100