use new decorator pattern, lapack trtri

Author: asifzubair

pytensor/tensor/rewriting/linalg.py

@@ -20,7 +20,6 @@
     ExtractDiag,
     Eye,
     TensorVariable,
-    Tri,
     concatenate,
     diag,
     diagonal,
@@ -53,6 +52,7 @@
     Solve,
     SolveBase,
     SolveTriangular,
+    TriangularInv,
     _bilinear_solve_discrete_lyapunov,
     block_diag,
     cholesky,
@@ -1033,14 +1033,6 @@ def _find_triangular_op(var):
     if is_lower or is_upper:
         return (is_lower, is_upper)
 
-    if var.owner and isinstance(var.owner.op, Tri):
-        # The 'k' parameter of Tri determines the diagonal.
-        # k=0 is the main diagonal.
-        k = var.owner.op.k
-        if k == 0:
-            is_lower = var.owner.op.lower
-            return (is_lower, not is_lower)
-
     if var.owner and isinstance(var.owner.op, Blockwise):
         core_op = var.owner.op.core_op
         if isinstance(core_op, Cholesky):
@@ -1051,16 +1043,13 @@ def _find_triangular_op(var):
 
 @register_canonicalize
 @register_stabilize
-@node_rewriter([Blockwise])
+@node_rewriter([blockwise_of(MATRIX_INVERSE_OPS)])
 def rewrite_inv_to_triangular_solve(fgraph, node):
     """
     This rewrite takes advantage of the fact that the inverse of a triangular
     matrix can be computed more efficiently than the inverse of a general
     matrix by using a triangular solve instead of a general matrix inverse.
     """
-    core_op = node.op.core_op
-    if not isinstance(core_op, ALL_INVERSE_OPS):
-        return None
 
     A = node.inputs[0]
     triangular_info = _find_triangular_op(A)
@@ -1069,5 +1058,5 @@ def rewrite_inv_to_triangular_solve(fgraph, node):
 
     is_lower, is_upper = triangular_info
     if is_lower or is_upper:
-        I = pt.eye(A.shape[0], dtype=A.dtype)
-        return [solve_triangular(A, I, lower=is_lower)]
+        new_op = TriangularInv(lower=is_lower)
+        return [new_op(A)]

pytensor/tensor/slinalg.py

@@ -21,7 +21,7 @@
 from pytensor.tensor import math as ptm
 from pytensor.tensor.basic import as_tensor_variable, diagonal
 from pytensor.tensor.blockwise import Blockwise
-from pytensor.tensor.nlinalg import kron, matrix_dot
+from pytensor.tensor.nlinalg import MatrixInverse, kron, matrix_dot
 from pytensor.tensor.shape import reshape
 from pytensor.tensor.type import matrix, tensor, vector
 from pytensor.tensor.variable import TensorVariable
@@ -950,6 +950,30 @@ def solve_triangular(
     return cast(TensorVariable, ret)
 
 
+class TriangularInv(MatrixInverse):
+    """
+    Computes the inverse of a triangular matrix.
+    """
+
+    __props__ = ("lower",)
+
+    def __init__(self, lower=True):
+        self.lower = lower
+
+    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
+            )
+        z[0] = inv
+
+
 class Solve(SolveBase):
     """
     Solve a system of linear equations.

tests/tensor/rewriting/test_linalg.py

@@ -23,6 +23,7 @@
     MatrixInverse,
     MatrixPinv,
     SLogDet,
+    inv,
     matrix_inverse,
     svd,
 )
@@ -34,6 +35,7 @@
     Solve,
     SolveBase,
     SolveTriangular,
+    TriangularInv,
     cho_solve,
     cholesky,
     solve,
@@ -1060,3 +1062,44 @@ 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_op():
+    x = matrix("x")
+    f_lower = function([x], Blockwise(TriangularInv(lower=True))(x))
+    f_upper = function([x], Blockwise(TriangularInv(lower=False))(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
+    )
+
+    # Test upper
+    a = np.triu(np.random.rand(5, 5) + 0.1)
+    a_inv = f_upper(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
+    )
+
+
+def test_inv_to_triangular_inv_rewrite():
+    x = matrix("x")
+
+    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
+    )
+
+    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
+    )