Added rewrite for inv(inv(x)) -> x

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -539,3 +539,21 @@ def svd_uv_merge(fgraph, node):
                     or len(fgraph.clients[cl.outputs[2]]) > 0
                 ):
                     return [cl.outputs[1]]
+
+
+@register_canonicalize
+@register_stabilize
+@node_rewriter([Blockwise])
+def rewrite_inv_inv(fgraph, node):
+    valid_inverses = (MatrixInverse, MatrixPinv)
+    if not (isinstance(node.op.core_op, valid_inverses)):
+        return None
+
+    potential_inner_inv = node.inputs[0].owner
+    if not (
+        isinstance(potential_inner_inv.op, Blockwise)
+        and isinstance(node.op.core_op, valid_inverses)
+    ):
+        return None
+
+    return [potential_inner_inv.inputs[0]]

tests/tensor/rewriting/test_linalg.py

@@ -9,6 +9,7 @@
 from pytensor import function
 from pytensor import tensor as pt
 from pytensor.compile import get_default_mode
+from pytensor.compile.ops import DeepCopyOp
 from pytensor.configdefaults import config
 from pytensor.tensor import swapaxes
 from pytensor.tensor.blockwise import Blockwise
@@ -545,3 +546,15 @@ def test_svd_uv_merge():
             assert node.op.compute_uv
             svd_counter += 1
     assert svd_counter == 1
+
+
+def test_inv_inv_rewrite():
+    x = pt.matrix("a")
+    ii_x = pt.linalg.inv(pt.linalg.inv(x))
+    f_rewritten = function([x], ii_x, mode="FAST_RUN")
+    nodes = f_rewritten.maker.fgraph.apply_nodes
+
+    assert all(isinstance(node.op, DeepCopyOp) for node in nodes)
+
+    x_testing = np.random.rand(10, 10)
+    np.testing.assert_allclose(f_rewritten(x_testing), x_testing)