Skip to content

Comments

FIX inv_as_solve rewrite correctness for matmul inverse patterns#1895

Open
aman-coder03 wants to merge 5 commits intopymc-devs:mainfrom
aman-coder03:enh-transpose-solve-simplification
Open

FIX inv_as_solve rewrite correctness for matmul inverse patterns#1895
aman-coder03 wants to merge 5 commits intopymc-devs:mainfrom
aman-coder03:enh-transpose-solve-simplification

Conversation

@aman-coder03
Copy link

@aman-coder03 aman-coder03 commented Feb 19, 2026

Description

Closes #1894

This PR improves the correctness and canonicalization of linear algebra rewrites involving matrix inverse and matmul operations.
The existing inv_as_solve rewrite handled Dot patterns but did not correctly handle _matmul patterns. In particular A @ inv(B) must be rewritten as solve(B.T, A.T).T to preserve correct algebraic semantics
This PR extends the rewrite to handle _matmul and ensures mathematically correct transformations
After inverse-to-solve rewrites, expressions such as transpose(solve(transpose(A), transpose(B))) may appear in the graph

This PR introduces a canonicalization rewrite transpose(solve(transpose(A), transpose(B))) → solve(A, B)
This removes redundant transpose operations and improves graph simplicity

Added tests to verify
• Solve rewrite is applied
• Simplification rewrite works correctly
• Numerical correctness is preserved

Related Issue

  • Closes 1894
  • Related to 1893

Checklist

Type of change

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

@aman-coder03
Copy link
Author

@ricardoV94 can you please have a look at this PR?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

ENH: Simplify transpose(solve(transpose(A), transpose(B))) patterns

1 participant