ENH: Implement rewrite for inverse of triangular matrices

Author: Youssef-Naggar

pytensor/tensor/rewriting/linalg.py

@@ -821,6 +821,37 @@ def rewrite_inv_diag_to_diag_reciprocal(fgraph, node):
     return [eye_input / non_eye_input]
 
 
+@register_canonicalize
+@register_stabilize
+@node_rewriter([blockwise_of(MATRIX_INVERSE_OPS)])
+def rewrite_inv_triangular_to_solve(fgraph, node):
+    """
+    Rewrite `inv(A)` -> `solve_triangular(A, I)` when A is triangular.
+    """
+    # node is a Blockwise(MatrixInverse). The input to MatrixInverse is node.inputs[0]
+    inputs = node.inputs[0]
+
+    # Check for tags
+    is_lower = getattr(inputs.tag, "lower_triangular", False)
+    is_upper = getattr(inputs.tag, "upper_triangular", False)
+
+    if is_lower or is_upper:
+        # Create an identity matrix of the same size.
+        # Note: We use the last dimension for the size of the square matrix
+        n = inputs.shape[-1]
+
+        # Ensure the dtype matches the input dtype
+        identity = pt.eye(n, dtype=inputs.type.dtype)
+
+        # We replace the slow Inverse with the fast Triangular Solve.
+        # b_ndim=2 because identity is a 2D matrix.
+        inv_val = solve_triangular(inputs, identity, lower=is_lower, b_ndim=2)
+
+        return [inv_val]
+
+    return None
+
+
 @register_canonicalize
 @register_stabilize
 @node_rewriter([ExtractDiag])

tests/tensor/rewriting/test_linalg.py

@@ -10,6 +10,7 @@
 from pytensor import tensor as pt
 from pytensor.compile import get_default_mode
 from pytensor.configdefaults import config
+from pytensor.gradient import verify_grad
 from pytensor.graph import FunctionGraph, ancestors
 from pytensor.graph.rewriting.utils import rewrite_graph
 from pytensor.tensor import swapaxes
@@ -1128,3 +1129,32 @@ def solve_op_in_graph(graph):
     np.testing.assert_allclose(
         f(a_val, b_val), c_val, rtol=1e-7 if config.floatX == "float64" else 1e-5
     )
+
+
+def test_triangular_inv_rewrite_and_grad():
+    np.random.seed(42)
+    A_val = np.tril(np.random.rand(5, 5) + 5.0)
+
+    A = pt.dmatrix("A")
+    A.tag.lower_triangular = True
+
+    Z = pt.linalg.inv(A)
+
+    f = pytensor.function([A], Z)
+
+    assert not any(
+        isinstance(getattr(node.op, "core_op", node.op), MatrixInverse)
+        for node in f.maker.fgraph.toposort()
+    )
+    assert any(
+        isinstance(getattr(node.op, "core_op", node.op), SolveTriangular)
+        for node in f.maker.fgraph.toposort()
+    )
+
+    def func(x):
+        x = pt.as_tensor_variable(x)
+        x = pt.tril(x)
+        x.tag.lower_triangular = True
+        return pt.sum(pt.linalg.inv(x))
+
+    verify_grad(func, [A_val])