added function for rewriting cholesky(eye) -> eye

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -536,3 +536,39 @@ 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_cholesky_eye_to_eye(fgraph, node):
+    """
+     This rewrite takes advantage of the fact that the cholesky decomposition 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 Cholesky.
+
+    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
+    """
+    # Find whether cholesky op is being applied
+    if not isinstance(node.op.core_op, Cholesky):
+        return None
+
+    # Check whether input to Cholesky 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]

tests/tensor/rewriting/test_linalg.py

@@ -545,3 +545,30 @@ def test_svd_uv_merge():
             assert node.op.compute_uv
             svd_counter += 1
     assert svd_counter == 1
+
+
+def test_cholesky_eye_rewrite():
+    x = pt.eye(10)
+    x_mat = pt.matrix("x")
+    L = pt.linalg.cholesky(x)
+    L_mat = pt.linalg.cholesky(x_mat)
+    f_rewritten = function([], L, mode="FAST_RUN")
+    f_rewritten_mat = function([x_mat], L_mat, mode="FAST_RUN")
+    nodes = f_rewritten.maker.fgraph.apply_nodes
+    nodes_mat = f_rewritten_mat.maker.fgraph.apply_nodes
+
+    # Rewrite Test
+    assert not any(isinstance(node.op, Cholesky) for node in nodes)
+    assert any(isinstance(node.op, Cholesky) for node in nodes_mat)
+
+    # Value Test
+    x_test = np.eye(10)
+    L = np.linalg.cholesky(x_test)
+    rewritten_val = f_rewritten()
+
+    assert_allclose(
+        L,
+        rewritten_val,
+        atol=1e-3 if config.floatX == "float32" else 1e-8,
+        rtol=1e-3 if config.floatX == "float32" else 1e-8,
+    )