ENH: Extend linear algebra graph rewrites to cover additional inverse and matmul patterns

[aman-coder03]

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.

[ricardoV94]

We should confirm those cases are not already handled

[aman-coder03]

I tested the cases locally and confirmed that both:

pt.matmul(pt.linalg.inv(A), B)
and

pt.matmul(A, pt.linalg.inv(B))
are already rewritten into equivalent solve-based expressions. The resulting computation graphs contain Solve and transpose operations, and do not retain explicit MatrixInverse nodes.

It looks like these inverse–matmul patterns are already handled correctly. I will continue reviewing other inverse-related compositions to check if there are still cases where explicit inverse operations are not optimized.

[aman-coder03]

@ricardoV94 i have opened a PR related to Consolidate typing in pymc, kindly have a look pymc-devs/pymc#8094

[aman-coder03]

While testing, I noticed that some of these rewrites produce patterns like Transpose(Solve(Transpose(A), Transpose(B)))
which are mathematically correct but introduce additional transpose operations
would it make sense to open a separate issue to discuss adding a rewrite to simplify transpose solve transpose compositions when possible?

[ricardoV94]

Transpose are basically free (it's just a view on the array with different strides/shape). Nice if we cam simplify (we do rewrite a transpose of dot as the dot of transpose for example) but not super critical. Feel free to open an issue documenting it

[aman-coder03]

@ricardoV94 i have opened an issue and also linked the PR (#1895 )