Added rewrite for inv(eye) and removed (for now) orthonormal rewrites

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -572,6 +572,40 @@ def svd_uv_merge(fgraph, node):
                     return [cl.outputs[1]]
 
 
+@register_canonicalize
+@register_stabilize
+@node_rewriter([Blockwise])
+def rewrite_inv_eye_to_eye(fgraph, node):
+    """
+     This rewrite takes advantage of the fact that the inverse of an identity matrix is the matrix itself
+    The presence of an identity matrix is identified by checking whether we have k = 0 for an Eye Op inside an inverse op.
+    Parameters
+    ----------
+    fgraph: FunctionGraph
+        Function graph being optimized
+    node: Apply
+        Node of the function graph to be optimized
+    Returns
+    -------
+    list of Variable, optional
+        List of optimized variables, or None if no optimization was performed
+    """
+    valid_inverses = (MatrixInverse, MatrixPinv)
+    core_op = node.op.core_op
+    if not (isinstance(core_op, valid_inverses)):
+        return None
+
+    # Check whether input to inverse is Eye and the 1's are on main diagonal
+    eye_check = node.inputs[0]
+    if not (
+        eye_check.owner
+        and isinstance(eye_check.owner.op, Eye)
+        and getattr(eye_check.owner.inputs[-1], "data", -1).item() == 0
+    ):
+        return None
+    return [eye_check]
+
+
 @register_canonicalize
 @register_stabilize
 @node_rewriter([Blockwise])
@@ -631,40 +665,40 @@ def rewrite_inv_diag_to_diag_reciprocal(fgraph, node):
         return [eye_input / non_eye_input]
 
 
-@register_canonicalize
-@register_stabilize
-@node_rewriter([Blockwise])
-def rewrite_inv_for_orthonormal(fgraph, node):
-    """
-     This rewrite takes advantage of the fact that for an orthonormal matrix, the inverse is simply the transpose.
-     This function deals with orthonormal matrix arising from pt.linalg.svd decomposition (U, Vh) or arising from pt.linalg.qr
-
-    Parameters
-    ----------
-    fgraph: FunctionGraph
-        Function graph being optimized
-    node: Apply
-        Node of the function graph to be optimized
-
-    Returns
-    -------
-    list of Variable, optional
-        List of optimized variables, or None if no optimization was performed
-    """
-    # Dealing with orthonormal matrix from SVD
-    # Check if input to Inverse is coming from SVD
-    input_to_inv = node.inputs[0]
-    # Check if this input is coming from SVD with compute_uv = True
-    if not (
-        input_to_inv.owner
-        and isinstance(input_to_inv.owner.op, Blockwise)
-        and isinstance(input_to_inv.owner.op.core_op, SVD)
-        and input_to_inv.owner.op.core_op.compute_uv is True
-    ):
-        return None
-
-    # To make sure input is orthonormal, we have to check that its not S (output order of SVD is U, S, Vh, so S is index 1) (S matrix consists of singular values and it is not orthonormal)
-    if input_to_inv == input_to_inv.owner.outputs[1]:
-        return None
-
-    return [input_to_inv.T]
+# @register_canonicalize
+# @register_stabilize
+# @node_rewriter([Blockwise])
+# def rewrite_inv_for_orthonormal(fgraph, node):
+#     """
+#      This rewrite takes advantage of the fact that for an orthonormal matrix, the inverse is simply the transpose.
+#      This function deals with orthonormal matrix arising from pt.linalg.svd decomposition (U, Vh) or arising from pt.linalg.qr
+
+#     Parameters
+#     ----------
+#     fgraph: FunctionGraph
+#         Function graph being optimized
+#     node: Apply
+#         Node of the function graph to be optimized
+
+#     Returns
+#     -------
+#     list of Variable, optional
+#         List of optimized variables, or None if no optimization was performed
+#     """
+#     # Dealing with orthonormal matrix from SVD
+#     # Check if input to Inverse is coming from SVD
+#     input_to_inv = node.inputs[0]
+#     # Check if this input is coming from SVD with compute_uv = True
+#     if not (
+#         input_to_inv.owner
+#         and isinstance(input_to_inv.owner.op, Blockwise)
+#         and isinstance(input_to_inv.owner.op.core_op, SVD)
+#         and input_to_inv.owner.op.core_op.compute_uv is True
+#     ):
+#         return None
+
+#     # To make sure input is orthonormal, we have to check that its not S (output order of SVD is U, S, Vh, so S is index 1) (S matrix consists of singular values and it is not orthonormal)
+#     if input_to_inv == input_to_inv.owner.outputs[1]:
+#         return None
+
+#     return [input_to_inv.T]

tests/tensor/rewriting/test_linalg.py

@@ -559,6 +559,29 @@ def test_svd_uv_merge():
     assert svd_counter == 1
 
 
+def test_inv_eye_to_eye():
+    x = pt.eye(10)
+    x_inv = pt.linalg.inv(x)
+    f_rewritten = function([], x_inv, mode="FAST_RUN")
+    nodes = f_rewritten.maker.fgraph.apply_nodes
+
+    # Rewrite Test
+    valid_inverses = (MatrixInverse, MatrixPinv)
+    assert not any(isinstance(node.op, valid_inverses) for node in nodes)
+
+    # Value Test
+    x_test = np.eye(10)
+    x_inv_val = np.linalg.inv(x_test)
+    rewritten_val = f_rewritten()
+
+    assert_allclose(
+        x_inv_val,
+        rewritten_val,
+        atol=1e-3 if config.floatX == "float32" else 1e-8,
+        rtol=1e-3 if config.floatX == "float32" else 1e-8,
+    )
+
+
 @pytest.mark.parametrize(
     "shape",
     [(), (7,), (7, 7)],