changed check for inverse and solve

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -595,30 +595,31 @@ def _find_solve_with_eye(node):
 @register_stabilize
 @node_rewriter([Blockwise])
 def rewrite_inv_inv(fgraph, node):
-    print(f"NODE - {node}")
     valid_inverses = (MatrixInverse, MatrixPinv, Solve, SolveTriangular)
     valid_solves = (Solve, SolveTriangular)
     # Check if Solve has b = eye
-    solve_inv_check = False
-    if hasattr(node.op, "core_op") and isinstance(node.op.core_op, valid_solves):
-        solve_inv_check = _find_solve_with_eye(node)
+    inv_check = False
+    if hasattr(node.op, "core_op") and isinstance(node.op.core_op, valid_inverses):
+        inv_check = True
+        if isinstance(node.op.core_op, valid_solves):
+            inv_check = _find_solve_with_eye(node)
 
-    if not solve_inv_check:
-        return None
-
-    if not (isinstance(node.op.core_op, valid_inverses)):
+    if not inv_check:
         return None
 
     potential_inner_inv = node.inputs[0].owner
     if potential_inner_inv is None or potential_inner_inv.op is None:
         return None
     # Check if its an inner solve as well, does that have b = eye
-    solve_inv_check = False
+    inv_check = False
     if hasattr(potential_inner_inv.op, "core_op") and isinstance(
-        potential_inner_inv.op.core_op, valid_solves
+        potential_inner_inv.op.core_op, valid_inverses
     ):
-        solve_inv_check = _find_solve_with_eye(potential_inner_inv)
-    if not solve_inv_check:
+        inv_check = True
+        if isinstance(potential_inner_inv.op.core_op, valid_solves):
+            inv_check = _find_solve_with_eye(potential_inner_inv)
+
+    if not inv_check:
         return None
 
     if not (

tests/tensor/rewriting/test_linalg.py

@@ -570,6 +570,7 @@ def get_pt_function(x, op_name):
     op1 = get_pt_function(x, inv_op_1)
     op2 = get_pt_function(op1, inv_op_2)
     f_rewritten = function([x], op2, mode="FAST_RUN")
+    print(f_rewritten.dprint())
     nodes = f_rewritten.maker.fgraph.apply_nodes
 
     valid_inverses = (MatrixInverse, MatrixPinv, Solve, SolveTriangular)