fixed rewrite for solve

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -36,6 +36,7 @@
     Cholesky,
     Solve,
     SolveBase,
+    SolveTriangular,
     block_diag,
     cholesky,
     solve,
@@ -542,28 +543,46 @@ def svd_uv_merge(fgraph, node):
 
 
 def _find_solve_with_eye(node):
+    valid_solves = (Solve, SolveTriangular)
     # First, we look for the solve op
-    if not (isinstance(node.op, Blockwise) and isinstance(node.op.core_op, Solve)):
-        return None
-
+    if not (
+        isinstance(node.op, Blockwise) and isinstance(node.op.core_op, valid_solves)
+    ):
+        return False
     # Check whether second input to solve is Eye
-    solve_inputs = node.owner.inputs
+    solve_inputs = node.inputs
     potential_eye_input = solve_inputs[1].owner
-    if not (isinstance(potential_eye_input.op, Eye)):
-        return None
+    if not (potential_eye_input and isinstance(potential_eye_input.op, Eye)):
+        return False
 
-    return node
+    return True
 
 
 @register_canonicalize
 @register_stabilize
 @node_rewriter([Blockwise])
 def rewrite_inv_inv(fgraph, node):
-    valid_inverses = (MatrixInverse, MatrixPinv)
+    valid_inverses = (MatrixInverse, MatrixPinv, Solve, SolveTriangular)
+    valid_solves = (Solve, SolveTriangular)
+    # Check if Solve has b = eye
+    solve_inv_check = True
+    if isinstance(node.op.core_op, valid_solves):
+        solve_inv_check = _find_solve_with_eye(node)
+
+    if not solve_inv_check:
+        return None
+
     if not (isinstance(node.op.core_op, valid_inverses)):
         return None
 
     potential_inner_inv = node.inputs[0].owner
+    # Check if its an inner solve as well, does that have b = eye
+    solve_inv_check = True
+    if isinstance(potential_inner_inv.op.core_op, valid_solves):
+        solve_inv_check = _find_solve_with_eye(potential_inner_inv)
+    if not solve_inv_check:
+        return None
+
     if not (
         potential_inner_inv
         and isinstance(potential_inner_inv.op, Blockwise)

tests/tensor/rewriting/test_linalg.py

@@ -547,17 +547,59 @@ def test_svd_uv_merge():
     assert svd_counter == 1
 
 
-@pytest.mark.parametrize(
-    "inv_op", [pt.linalg.inv, pt.linalg.pinv], ids=["MatrixInverse", "MatrixPinv"]
-)
-def test_inv_inv_rewrite(inv_op):
+def test_inv_inv_rewrite():
     x = pt.matrix("a")
-    ii_x = inv_op(inv_op(x))
-    f_rewritten = function([x], ii_x, mode="FAST_RUN")
-    nodes = f_rewritten.maker.fgraph.apply_nodes
-
+    inv_ops = [pt.linalg.inv, pt.linalg.pinv]
+    solve_ops = [pt.linalg.solve, pt.linalg.solve_triangular]
     valid_inverses = (MatrixInverse, MatrixPinv)
-    assert not any(isinstance(node.op, valid_inverses) for node in nodes)
-
-    x_testing = np.random.rand(10, 10)
-    np.testing.assert_allclose(f_rewritten(x_testing), x_testing)
+    valid_solves = (Solve, SolveTriangular)
+    all_valid = (MatrixInverse, MatrixPinv, Solve, SolveTriangular)
+    # inv(inv)
+    for inv_op in inv_ops:
+        for inv_op_2 in inv_ops:
+            ii_x = inv_op(inv_op_2(x))
+            f_rewritten = function([x], ii_x, mode="FAST_RUN")
+            nodes = f_rewritten.maker.fgraph.apply_nodes
+
+            assert not any(isinstance(node.op, valid_inverses) for node in nodes)
+
+            x_testing = np.random.rand(10, 10)
+            np.testing.assert_allclose(f_rewritten(x_testing), x_testing)
+    # solve(solve)
+    b_eye = pt.eye(10)
+    for solve_op in solve_ops:
+        for solve_op_2 in solve_ops:
+            ss_x = solve_op(solve_op_2(x, b_eye), b_eye)
+            with pytensor.config.change_flags(optimizer_verbose=True):
+                f_rewritten = function([x], ss_x, mode="FAST_RUN")
+            nodes = f_rewritten.maker.fgraph.apply_nodes
+
+            assert not any(isinstance(node.op, valid_solves) for node in nodes)
+
+            x_testing = np.random.rand(10, 10)
+            np.testing.assert_allclose(f_rewritten(x_testing), x_testing)
+
+    # inv(solve)
+    for inv_op in inv_ops:
+        for solve_op in solve_ops:
+            is_x = inv_op(solve_op(x, b_eye))
+            with pytensor.config.change_flags(optimizer_verbose=True):
+                f_rewritten = function([x], is_x, mode="FAST_RUN")
+            nodes = f_rewritten.maker.fgraph.apply_nodes
+            assert not any(isinstance(node.op, all_valid) for node in nodes)
+
+            x_testing = np.random.rand(10, 10)
+            np.testing.assert_allclose(f_rewritten(x_testing), x_testing)
+
+    # solve(inv)
+    for solve_op in solve_ops:
+        for inv_op in inv_ops:
+            si_x = solve_op(inv_op(x), b_eye)
+            with pytensor.config.change_flags(optimizer_verbose=True):
+                f_rewritten = function([x], si_x, mode="FAST_RUN")
+            nodes = f_rewritten.maker.fgraph.apply_nodes
+
+            assert not any(isinstance(node.op, all_valid) for node in nodes)
+
+            x_testing = np.random.rand(10, 10)
+            np.testing.assert_allclose(f_rewritten(x_testing), x_testing)