address review comments;

add other conditions to trigger rewrite
enhance TriInv Op
add tests

Author: asifzubair

pytensor/tensor/rewriting/linalg.py

@@ -8,18 +8,22 @@
 from pytensor import tensor as pt
 from pytensor.compile import optdb
 from pytensor.graph import Apply, FunctionGraph
+from pytensor.graph.basic import Constant
 from pytensor.graph.rewriting.basic import (
     copy_stack_trace,
     dfs_rewriter,
     node_rewriter,
 )
 from pytensor.graph.rewriting.unify import OpPattern
 from pytensor.scalar.basic import Abs, Log, Mul, Sign
+from pytensor.scalar.basic import Mul as ScalarMul
+from pytensor.scalar.basic import Sub as ScalarSub
 from pytensor.tensor.basic import (
     AllocDiag,
     ExtractDiag,
     Eye,
     TensorVariable,
+    Tri,
     concatenate,
     diag,
     diagonal,
@@ -46,9 +50,12 @@
 )
 from pytensor.tensor.rewriting.blockwise import blockwise_of
 from pytensor.tensor.slinalg import (
+    LU,
+    QR,
     BlockDiagonal,
     Cholesky,
     CholeskySolve,
+    LUFactor,
     Solve,
     SolveBase,
     SolveTriangular,
@@ -1026,17 +1033,69 @@ def _find_triangular_op(var):
 
     Returns a tuple (is_lower, is_upper) if triangular, otherwise None.
     """
-
+    # Case 1: Check for an explicit tag
     is_lower = getattr(var.tag, "lower_triangular", False)
     is_upper = getattr(var.tag, "upper_triangular", False)
-
     if is_lower or is_upper:
         return (is_lower, is_upper)
 
-    if var.owner and isinstance(var.owner.op, Blockwise):
-        core_op = var.owner.op.core_op
-        if isinstance(core_op, Cholesky):
-            return (core_op.lower, not core_op.lower)
+    if not var.owner:
+        return None
+
+    op = var.owner.op
+    core_op = op.core_op if isinstance(op, Blockwise) else op
+
+    # Case 2: Check for direct creator Ops
+    if isinstance(core_op, Cholesky):
+        return (core_op.lower, not core_op.lower)
+
+    if isinstance(core_op, LU | LUFactor):
+        if var.owner.outputs[1] == var:
+            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)
+
+    # pt.tri will get constant folded so no point re-writing ?
+    # if isinstance(core_op, Tri):
+    #     k_node = var.owner.inputs[2]
+    #     if isinstance(k_node, Constant) and k_node.data == 0:
+    #         print('re-writing ... ')
+    #         return (True, False)
+
+    # Case 3: tril/triu patterns which are implemented as Mul
+    if isinstance(core_op, Elemwise) and isinstance(core_op.scalar_op, ScalarMul):
+        other_inp = next(
+            (i for i in var.owner.inputs if i != var.owner.inputs[0]), None
+        )
+
+        if other_inp is not None and other_inp.owner:
+            # Check for tril pattern: Mul(x, Tri(...))
+            if isinstance(other_inp.owner.op, Tri):
+                k_node = other_inp.owner.inputs[2]
+                if isinstance(k_node, Constant) and k_node.data == 0:
+                    return (True, False)  # It's tril
+
+            # Check for triu pattern: Mul(x, Sub(1, Tri(k=-1)))
+            sub_op = other_inp.owner.op
+            if isinstance(sub_op, Elemwise) and isinstance(sub_op.scalar_op, ScalarSub):
+                sub_inputs = other_inp.owner.inputs
+                const_one = next(
+                    (i for i in sub_inputs if isinstance(i, Constant) and i.data == 1),
+                    None,
+                )
+                tri_inp = next(
+                    (i for i in sub_inputs if i.owner and isinstance(i.owner.op, Tri)),
+                    None,
+                )
+
+                if const_one is not None and tri_inp is not None:
+                    k_node = tri_inp.owner.inputs[2]
+                    if isinstance(k_node, Constant) and k_node.data == -1:
+                        return (False, True)  # It's triu
 
     return None
 
@@ -1059,4 +1118,6 @@ def rewrite_inv_to_triangular_solve(fgraph, node):
     is_lower, is_upper = triangular_info
     if is_lower or is_upper:
         new_op = TriangularInv(lower=is_lower)
-        return [new_op(A)]
+        new_inv = new_op(A)
+        copy_stack_trace(node.outputs[0], new_inv)
+        return [new_inv]

pytensor/tensor/slinalg.py

@@ -1021,24 +1021,46 @@ class TriangularInv(MatrixInverse):
     Computes the inverse of a triangular matrix.
     """
 
-    __props__ = ("lower",)
+    __props__ = ("lower", "on_error", "overwrite_a")
 
-    def __init__(self, lower=True):
+    def __init__(self, lower=True, on_error="raise", overwrite_a=False):
         self.lower = lower
+        if on_error not in ("raise", "nan"):
+            raise ValueError('on_error must be one of "raise" or "nan"')
+        self.on_error = on_error
+        self.overwrite_a = overwrite_a
+
+        if self.overwrite_a:
+            self.destroy_map = {0: [0]}
 
     def perform(self, node, inputs, outputs):
         (x,) = inputs
         (z,) = outputs
         (dtrtri,) = get_lapack_funcs(("trtri",), (x,))
         inv, info = dtrtri(x, lower=self.lower, overwrite_c=True)
-        if info > 0:
-            raise np.linalg.LinAlgError("Singular matrix")
-        elif info < 0:
-            raise ValueError(
-                "illegal value in %d-th argument of internal trtri" % -info
-            )
+        if info != 0:
+            if self.on_error == "nan":
+                z[0] = np.full_like(x, np.nan)
+                return
+            elif info > 0:
+                raise np.linalg.LinAlgError("Singular matrix")
+            elif info < 0:
+                raise ValueError(
+                    f"illegal value in {-info}-th argument of internal trtri"
+                )
         z[0] = inv
 
+    def inplace_on_inputs(self, allowed_inplace_inputs: list[int]) -> "Op":
+        """
+        Allows this Op to overwrite its input buffer with its output.
+        """
+        if not allowed_inplace_inputs:
+            return self
+
+        new_props = self._props_dict()
+        new_props["overwrite_a"] = True
+        return type(self)(**new_props)
+
 
 class Solve(SolveBase):
     """

tests/tensor/rewriting/test_linalg.py

@@ -11,8 +11,10 @@
 from pytensor.compile import get_default_mode
 from pytensor.configdefaults import config
 from pytensor.graph import ancestors
+from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.utils import rewrite_graph
 from pytensor.tensor import swapaxes
+from pytensor.tensor.basic import tril, triu
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import DimShuffle
 from pytensor.tensor.math import dot, matmul
@@ -38,6 +40,8 @@
     TriangularInv,
     cho_solve,
     cholesky,
+    lu,
+    qr,
     solve,
     solve_triangular,
 )
@@ -1064,42 +1068,121 @@ def solve_op_in_graph(graph):
     )
 
 
-def test_triangular_inv_op():
+@pytest.mark.parametrize("lower", [True, False])
+def test_triangular_inv_op(lower):
+    """Tests the TriangularInv Op directly."""
     x = matrix("x")
-    f_lower = function([x], Blockwise(TriangularInv(lower=True))(x))
-    f_upper = function([x], Blockwise(TriangularInv(lower=False))(x))
+    f = function([x], TriangularInv(lower=lower)(x))
 
-    # Test lower
-    a = np.tril(np.random.rand(5, 5) + 0.1)
-    a_inv = f_lower(a)
-    expected_inv = np.linalg.inv(a)
-    np.testing.assert_allclose(
-        np.tril(a_inv), np.tril(expected_inv), rtol=1e-5, atol=1e-7
-    )
+    if lower:
+        a = np.tril(np.random.rand(5, 5) + 0.1)
+    else:
+        a = np.triu(np.random.rand(5, 5) + 0.1)
 
-    # Test upper
-    a = np.triu(np.random.rand(5, 5) + 0.1)
-    a_inv = f_upper(a)
+    a_inv = f(a)
     expected_inv = np.linalg.inv(a)
-    np.testing.assert_allclose(
-        np.triu(a_inv), np.triu(expected_inv), rtol=1e-5, atol=1e-7
+
+    # Clean the NumPy result before comparing.
+    if lower:
+        expected_inv = np.tril(expected_inv)
+    else:
+        expected_inv = np.triu(expected_inv)
+
+    # The inverse of a triangular matrix is also triangular.
+    # We should check the full matrix, not just a part of it.
+    assert_allclose(
+        a_inv, expected_inv, rtol=1e-7 if config.floatX == "float64" else 1e-5
     )
 
 
-def test_inv_to_triangular_inv_rewrite():
+def test_triangular_inv_op_nan_on_error():
+    """
+    Tests the `on_error='nan'` functionality of the TriangularInv Op.
+    """
     x = matrix("x")
+    f_nan = function([x], TriangularInv(on_error="nan")(x))
+
+    # Create a singular triangular matrix (zero on the diagonal)
+    a_singular = np.tril(np.random.rand(5, 5))
+    a_singular[2, 2] = 0
 
-    x_chol = cholesky(x)
-    y_chol = inv(x_chol)
-    f_chol = function([x], y_chol)
-    assert any(
-        isinstance(node.op, TriangularInv)
-        or (hasattr(node.op, "core_op") and isinstance(node.op.core_op, TriangularInv))
-        for node in f_chol.maker.fgraph.apply_nodes
+    res = f_nan(a_singular)
+    assert np.all(np.isnan(res))
+
+
+def _check_op_in_graph(fgraph, op_type, present=True):
+    """Helper to check if an Op is in a graph."""
+
+    # We use type() instead of isinstance() to avoid matching subclasses
+    # (e.g., finding TriangularInv when we're looking for MatrixInverse).
+    found = any(
+        type(node.op) is op_type
+        or (hasattr(node.op, "core_op") and type(node.op.core_op) is op_type)
+        for node in fgraph.apply_nodes
     )
+    if present:
+        assert found, f"{op_type.__name__} not found in graph"
+    else:
+        assert not found, f"{op_type.__name__} unexpectedly found in graph"
+
+
+rewrite_cases = {
+    "tril": (
+        lambda x: tril(x),
+        lambda a: np.tril(a),
+    ),
+    "triu": (
+        lambda x: triu(x),
+        lambda a: np.triu(a),
+    ),
+    "cholesky": (
+        lambda x: cholesky(x),
+        lambda a: np.linalg.cholesky(a),
+    ),
+    "lu_L": (
+        lambda x: lu(x)[1],
+        lambda a: scipy.linalg.lu(a)[1],
+    ),
+    "lu_U": (
+        lambda x: lu(x)[2],
+        lambda a: scipy.linalg.lu(a)[2],
+    ),
+    "qr_R": (
+        lambda x: qr(x)[1],
+        lambda a: np.linalg.qr(a)[1],
+    ),
+}
+
+
+@pytest.mark.parametrize("case", rewrite_cases.keys())
+def test_inv_to_triangular_inv_rewrite(case):
+    """
+    Tests the rewrite of inv(triangular) -> TriangularInv.
+    """
+    x = matrix("x")
+    build_tri, _ = rewrite_cases[case]
+    x_tri = build_tri(x)
+    y_inv = inv(x_tri)
 
+    # Check graph BEFORE compilation
+    pre_compile_fgraph = FunctionGraph([x], [y_inv], clone=False)
+    _check_op_in_graph(pre_compile_fgraph, MatrixInverse, present=True)
+    _check_op_in_graph(pre_compile_fgraph, TriangularInv, present=False)
+
+    # Trigger the rewrite
+    f = function([x], y_inv)
+
+    # Check graph AFTER compilation
+    post_compile_fgraph = f.maker.fgraph
+    _check_op_in_graph(post_compile_fgraph, TriangularInv, present=True)
+    _check_op_in_graph(post_compile_fgraph, MatrixInverse, present=False)
+
+    # Check numerical correctness
     a = np.random.rand(5, 5)
-    a = np.dot(a, a.T) + np.eye(5) * 0.1  # ensure positive definite
-    np.testing.assert_allclose(
-        f_chol(a), np.linalg.inv(np.linalg.cholesky(a)), rtol=1e-5, atol=1e-7
+    a = np.dot(a, a.T) + np.eye(5)  # Make positive definite for Cholesky
+    pytensor_result = f(a)
+    _, numpy_tri_func = rewrite_cases[case]
+    numpy_result = np.linalg.inv(numpy_tri_func(a))
+    assert_allclose(
+        pytensor_result, numpy_result, rtol=1e-7 if config.floatX == "float64" else 1e-5
     )