added rewrite for inverse of triangular matrix

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -1038,3 +1038,22 @@ def det_triangular_to_prod_diag(fgraph, node):
         return [det_val]
 
     return None
+
+
+@register_canonicalize
+@register_stabilize
+@node_rewriter([Blockwise])
+def rewrite_inv_triangular_to_solve_triangular(fgraph, node):
+    core_op = node.op.core_op
+    if not (isinstance(core_op, ALL_INVERSE_OPS)):
+        return None
+
+    inputs = node.inputs[0]
+    triangular_check = _find_triangular_from_cholesky(inputs)
+
+    if triangular_check:
+        valid_eye = pt.eye(inputs.shape[-1])
+        inv_val = solve_triangular(inputs, valid_eye, lower=True)
+        return [inv_val]
+
+    return None

tests/tensor/rewriting/test_linalg.py

@@ -1009,7 +1009,7 @@ def test_det_triangular():
     nodes = f_rewritten.maker.fgraph.apply_nodes
     assert not any(isinstance(node.op, Det) for node in nodes)
 
-    # Numeric test
+    # Numeric Test
     x_test = np.random.rand(10, 10).astype(config.floatX)
     x_psd = np.dot(x_test, x_test.T)
     x_triangular = np.linalg.cholesky(x_psd)
@@ -1029,3 +1029,31 @@ def test_det_triangular():
     f_rewritten = function([y], z, mode="FAST_RUN")
     nodes = f_rewritten.maker.fgraph.apply_nodes
     assert any(isinstance(node.op, Det) for node in nodes)
+
+
+@pytest.mark.parametrize("inv_op", ["inv", "pinv"])
+def test_inv_triangular(inv_op):
+    x = pt.matrix("x")
+    x_triangular = pt.linalg.cholesky(x)
+    z = get_pt_function(x_triangular, inv_op)
+
+    # Rewrite Test
+    f_rewritten = function([x], z, mode="FAST_RUN")
+    nodes = f_rewritten.maker.fgraph.apply_nodes
+
+    valid_inverses = (MatrixInverse, MatrixPinv)
+    assert not any(isinstance(node.op, valid_inverses) for node in nodes)
+
+    # Numeric Test
+    x_test = np.random.rand(10, 10).astype(config.floatX)
+    x_psd = np.dot(x_test, x_test.T)
+    x_triangular = np.linalg.cholesky(x_psd)
+    inv_val = np.linalg.inv(x_triangular)
+    rewritten_val = f_rewritten(x_psd)
+
+    assert_allclose(
+        inv_val,
+        rewritten_val,
+        atol=1e-3 if config.floatX == "float32" else 1e-8,
+        rtol=1e-3 if config.floatX == "float32" else 1e-8,
+    )