Rewrite inverse for triangular matrix #1612

asifzubair · 2025-09-16T02:52:50Z

Description

We add a rewrite for matrix inversion when the matrix is triangular.

We check three conditions:

If there is a tag which is upper/lower triangular
If the Op is Tri
If the Op is Cholesky

Related Issue

Closes #
Related to ENH: Implement COLA library rewrites for linear algebra functions #573

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

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

asifzubair · 2025-09-16T02:54:32Z

Hi @jessegrabowski , I haven't added a test yet, but if this approach is valid, I can add one. Please let me know. Thank you! 🙏

CC: @theorashid , @ColtAllen

jessegrabowski · 2025-09-17T11:06:21Z

Can you post some timings showing that this is advantageous? inv(A) and solve(A, Eye) are basically the same thing. I recognize that there is some advantage because you're using the specialized solve, but I'd like to see what it gives us.

Can you also time solve_triangular(A, Eye) against directly using DTRTRI? That's the specialized LAPACK routine for the inverse of a triangular matrix (as opposed to a solve).

asifzubair · 2025-09-19T14:15:03Z

Thank you, @jessegrabowski . I'd be interested in seeing those numbers as well. Let me do that study and report back. Thanks, again 🙏

asifzubair · 2025-09-22T00:49:19Z

Hi @jessegrabowski , I did the study that was suggested. I used the underlying perform methods to do the benchmarks for inv(), solve() and solve_triangular().

Size	inv()	solve()	solve_tri()	dtrtri()	(inv/solve_tri)	(solve_tri/dtrtri)
50	0.01832s	0.00713s	0.01206s	0.00096s	1.52x	12.59x
100	0.02006s	0.01739s	0.00905s	0.00358s	2.22x	2.53x
250	0.09416s	0.10159s	0.04149s	0.06596s	2.27x	0.63x
500	0.43439s	0.57095s	0.12388s	0.07183s	3.51x	1.72x
750	1.07380s	1.16674s	0.46847s	0.18619s	2.29x	2.52x
1000	2.64248s	2.70513s	1.07507s	0.40498s	2.46x	2.65x
2000	22.21642s	27.74535s	8.88119s	2.54412s	2.50x	3.49x

On average, we see ~2X improvement when using solve_triangular() and then a further ~2X improvement when using dtrtri(). Please let me know your thoughts. Thank you 🙏

Click to view benchmarking code

import timeit
import numpy as np
import scipy.linalg
from scipy.linalg.lapack import dtrtri

matrix_sizes = [50, 100, 250, 500, 750, 1000, 2000]
n_repeats = 100

results = {}
for size in matrix_sizes:
    print(f"Running for size {size}x{size}...")
    
    A_tril = np.tril(np.random.rand(size, size))
    A_tril[np.diag_indices(size)] += 1.0
    I = np.eye(size)

    t_inv = timeit.timeit(lambda: np.linalg.inv(A_tril), number=n_repeats)
    t_solve = timeit.timeit(lambda: np.linalg.solve(A_tril, I), number=n_repeats)    
    t_solve_tri = timeit.timeit(
        lambda: scipy.linalg.solve_triangular(A_tril, I, lower=True),
        number=n_repeats
    )
    A_fortran = np.asfortranarray(A_tril)
    t_dtrtri = timeit.timeit(
        lambda: dtrtri(A_fortran, lower=1),
        number=n_repeats
    )
    
    results[size] = {
        "inv": t_inv,
        "solve": t_solve,
        "solve_triangular": t_solve_tri,
        "dtrtri": t_dtrtri,
        "inv_div_solve_tri": t_inv / t_solve_tri if t_solve_tri > 0 else 0,
        "solve_tri_div_dtrtri": t_solve_tri / t_dtrtri if t_dtrtri > 0 else 0
    }

jessegrabowski · 2025-09-22T00:56:18Z

That's really awesome! Thanks for doing this study.

Given these results, my suggestion would be to make a TriangularInv Op that subclasses from Inv. In the perform method, use trtri directly (use get_lapack_function to get it over directly importing it so you get the right one based on the dtype of the inputs). We don't have to expose it to users, but you can use it in rewrites.

We can also then add a rewrite that changes TriangularInv(x) @ b to solve_triangular(x, b). But that's another PR for another day :)

codecov · 2025-09-22T15:13:54Z

Codecov Report

❌ Patch coverage is 74.73684% with 24 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.62%. Comparing base (ab5037e) to head (14fec15).

Files with missing lines	Patch %	Lines
pytensor/tensor/rewriting/linalg.py	78.18%	4 Missing and 8 partials ⚠️
pytensor/tensor/slinalg.py	69.23%	10 Missing and 2 partials ⚠️

❌ Your patch check has failed because the patch coverage (74.73%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1612      +/-   ##
==========================================
- Coverage   81.64%   81.62%   -0.02%     
==========================================
  Files         244      244              
  Lines       53590    53683      +93     
  Branches     9438     9464      +26     
==========================================
+ Hits        43752    43821      +69     
- Misses       7356     7370      +14     
- Partials     2482     2492      +10

Files with missing lines	Coverage Δ
pytensor/tensor/nlinalg.py	`94.45% <100.00%> (ø)`
pytensor/tensor/rewriting/linalg.py	`89.50% <78.18%> (-1.59%)`	⬇️
pytensor/tensor/slinalg.py	`90.47% <69.23%> (-0.94%)`	⬇️

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

ricardoV94 · 2025-09-23T09:21:55Z

pytensor/tensor/rewriting/linalg.py

+    core_op = node.op.core_op
+    if not isinstance(core_op, ALL_INVERSE_OPS):
+        return None


We've merged some changes recently, so you can basically put this in the tracks: #1594

Thank you, @ricardoV94 . Just checking, can I eliminate this conditional in favor of this decorator:

@node_rewriter([blockwise_of(MATRIX_INVERSE_OPS)])

asifzubair · 2025-09-29T04:04:31Z

Hi @jessegrabowski , could you please review when you have the time ? I do checks similar to this PR, except, of course, we use the lapack solver. I'm also happy to include other operations like det and eig into this PR, if you recommend.

Also, if helpful, I can start a separate issue to track TriangularInv(x) @ b.

Please let me know. Thanks!

jessegrabowski

Thanks for the ping! This is looking really amazing, and it's getting very close 🥳

pytensor/tensor/rewriting/linalg.py

jessegrabowski · 2025-10-01T02:28:29Z

pytensor/tensor/rewriting/linalg.py

+    is_upper = getattr(var.tag, "upper_triangular", False)
+
+    if is_lower or is_upper:
+        return (is_lower, is_upper)


just returning one should be sufficient?

I thought that this gives some flexibility for returning diagonal etc. (is_upper = True and is_lower = True). No strong opinion though.

My preference would be for a separate is_diagonal helper, and to simplify this to just one return type.

Long term, I'm hoping we will have op-by-op matrix type inference, so these checks will be much easier.

jessegrabowski · 2025-10-01T02:30:38Z

pytensor/tensor/rewriting/linalg.py

+    is_lower, is_upper = triangular_info
+    if is_lower or is_upper:
+        new_op = TriangularInv(lower=is_lower)
+        return [new_op(A)]


Need to copy_stack_trace

jessegrabowski · 2025-10-01T02:33:07Z

pytensor/tensor/slinalg.py

+        if info > 0:
+            raise np.linalg.LinAlgError("Singular matrix")
+        elif info < 0:
+            raise ValueError(
+                "illegal value in %d-th argument of internal trtri" % -info


I would prefer if we return np.full_like(x, np.nan) when the algorithm fails. This is what jax does -- its very frustrating to have iterative algorithms (like mcmc/sgd) totally stop because of an unstable linalg operation.

Check how the Cholesky Op handles it.

pytensor/tensor/slinalg.py

jessegrabowski · 2025-10-01T02:39:51Z

tests/tensor/rewriting/test_linalg.py

+    x_chol = cholesky(x)
+    y_chol = inv(x_chol)
+    f_chol = function([x], y_chol)
+    assert any(


Also check that there is a regular Inv Ops in the graph before compiling, but not after.

jessegrabowski · 2025-10-01T02:41:14Z

Also, if helpful, I can start a separate issue to track TriangularInv(x) @ b.

Sure, feel free to open an issue. If you prefer to wait for this to be merged, that's fine too.

asifzubair · 2025-10-05T20:55:32Z

Thanks for the ping! This is looking really amazing, and it's getting very close 🥳

Thank you, @jessegrabowski ! 🙏 I took a stab at the comments. Notably, I added triu and tril checks. However, when I tried adding the pt.tri check, it seems to get const folded so doesn't trigger the re-write. Happy to be corrected.

Please let me know your thoughts. Thanks, again! 🙏

jessegrabowski · 2025-10-05T23:02:43Z

So you know, you can run mypy locally from inside the pytensor project folder with python scripts/run_mypy.py --verbose. You then get a huge horrible read-out that you have to sift though for the list of recent failures. I strongly recommend clearing the terminal before you run it.

jessegrabowski · 2025-10-05T23:12:42Z

For the failed float32 test, make sure you set the atol and rtol much more relaxed when config.floatX is float32. Check the other tests to see what we do. We need to think of a better way to test linalg routines at half-precision...

asifzubair · 2025-10-05T23:57:48Z

Thank you, @jessegrabowski . Yes, sorry, I realized belatedly about the mypy check

I ran it locally and it seems that I'm fighting this error:

[pytensor/tensor/slinalg.py]
pytensor/tensor/slinalg.py:958: error [assignment]: Incompatible types in assignment (expression has type "tuple[str, str, str]", base class "MatrixInverse" defined the type as "tuple[()]")

Should I redefine the __props__ type in TriangularInv ?

For the test failure, I had tried following the idiom seen elsewhere:

    np.testing.assert_allclose(
        f(a_val, b_val), c_val, rtol=1e-7 if config.floatX == "float64" else 1e-5
    )

But perhaps I need to loosen up the requirement (1e-4) for my test. Would that be okay ?

I also realized some of my tests are in the wrong location, the TriangularInv Op class tests should be in tests/tensor/test_slinalg.py and not in rewriting. Sorry about that as well, I'll move them.

Also, a heads up, I have to do some convoluted checks for re-writing triu / tril . I've tested that these do work, in that, they trigger the re-write and return a correct answer, however, I'm not sure if there is an easier way to do check. Please let me know if there is. Thank you. 🙏

jessegrabowski

This PR is really amazing, we're super close. Sorry for being slow, let's try to get this merged in the next couple days!

jessegrabowski · 2025-10-14T01:47:22Z

pytensor/tensor/rewriting/linalg.py

+    is_upper = getattr(var.tag, "upper_triangular", False)
+
+    if is_lower or is_upper:
+        return (is_lower, is_upper)


My preference would be for a separate is_diagonal helper, and to simplify this to just one return type.

Long term, I'm hoping we will have op-by-op matrix type inference, so these checks will be much easier.

pytensor/tensor/rewriting/linalg.py

tests/tensor/test_slinalg.py

jessegrabowski · 2025-10-14T02:08:15Z

tests/tensor/test_slinalg.py

+
+
+@pytest.mark.parametrize("lower", [True, False])
+def test_triangular_inv_op(lower):


You also need to test that the overwriting is working correctly. To do this you need to use pytensor.In(x, mutable=overwrite_a) in pytensor.function. Otherwise input variables are always treated as immutable and never overwritten. Check here for a rough example to follow

Hi @jessegrabowski , I've tried implementing this test, but not feeling very confident. Could you please review and let me know your thoughts ? Thank you! 🙏

add other conditions to trigger rewrite enhance TriInv Op add tests

jessegrabowski

Double check the is_triangular check for the LU/QR cases, then I think this is done! Really great work!

jessegrabowski · 2025-10-26T18:13:15Z

pytensor/tensor/rewriting/linalg.py

+            return (True, False)
+        if var.owner.outputs[2] == var:
+            return (False, True)
+
+    if isinstance(core_op, QR):
+        if var.owner.outputs[1] == var:
+            return (False, True)


These cases are still returning a tuple

pytensor/tensor/rewriting/linalg.py

asifzubair marked this pull request as draft September 16, 2025 02:53

ricardoV94 reviewed Sep 23, 2025

View reviewed changes

asifzubair force-pushed the azubair/enh-573-rewrite-inv-triangular branch from 7d7fcef to 34b3eb8 Compare September 26, 2025 04:25

asifzubair marked this pull request as ready for review September 29, 2025 03:57

asifzubair requested a review from ricardoV94 September 29, 2025 04:04

jessegrabowski requested changes Oct 1, 2025

View reviewed changes

asifzubair force-pushed the azubair/enh-573-rewrite-inv-triangular branch 2 times, most recently from d6e519e to 7f92bf6 Compare October 10, 2025 18:47

asifzubair requested a review from jessegrabowski October 10, 2025 18:48

jessegrabowski requested changes Oct 14, 2025

View reviewed changes

asifzubair force-pushed the azubair/enh-573-rewrite-inv-triangular branch from 7f92bf6 to 427008d Compare October 18, 2025 21:10

asifzubair added 7 commits October 26, 2025 10:08

add triangular rewrite

7c11584

use new decorator pattern, lapack trtri

a1e19e7

address review comments;

8d54703

add other conditions to trigger rewrite enhance TriInv Op add tests

fix tests & mypy issues

4c5e21d

fix mypy error, fix tests tol, move tests

c7af980

typo

208316d

review comments: overwrite_a test + tri rewrite test

14fec15

asifzubair force-pushed the azubair/enh-573-rewrite-inv-triangular branch from 427008d to 14fec15 Compare October 26, 2025 15:09

asifzubair requested a review from jessegrabowski October 26, 2025 15:09

jessegrabowski requested changes Oct 26, 2025

View reviewed changes



		@pytest.mark.parametrize("lower", [True, False])
		def test_triangular_inv_op(lower):

Rewrite inverse for triangular matrix #1612

Are you sure you want to change the base?

Rewrite inverse for triangular matrix #1612

Conversation

asifzubair commented Sep 16, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Checklist

Type of change

Uh oh!

asifzubair commented Sep 16, 2025

Uh oh!

jessegrabowski commented Sep 17, 2025

Uh oh!

asifzubair commented Sep 19, 2025

Uh oh!

asifzubair commented Sep 22, 2025

Uh oh!

jessegrabowski commented Sep 22, 2025

Uh oh!

codecov bot commented Sep 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

asifzubair commented Sep 29, 2025

Uh oh!

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jessegrabowski commented Oct 1, 2025

Uh oh!

asifzubair commented Oct 5, 2025

Uh oh!

jessegrabowski commented Oct 5, 2025

Uh oh!

jessegrabowski commented Oct 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

asifzubair commented Oct 5, 2025

Uh oh!

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

asifzubair commented Sep 16, 2025 •

edited by github-actions bot

Loading

codecov bot commented Sep 22, 2025 •

edited

Loading

jessegrabowski commented Oct 5, 2025 •

edited

Loading