minor changes; added test to not apply rewrite

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -617,7 +617,7 @@ def rewrite_inv_inv(fgraph, node):
 @register_canonicalize
 @register_stabilize
 @node_rewriter([Blockwise])
-def rewrite_cholesky_eye_to_eye(fgraph, node):
+def rewrite_remove_useless_cholesky(fgraph, node):
     """
      This rewrite takes advantage of the fact that the cholesky decomposition of an identity matrix is the matrix itself
 
@@ -640,14 +640,15 @@ def rewrite_cholesky_eye_to_eye(fgraph, node):
         return None
 
     # Check whether input to Cholesky is Eye and the 1's are on main diagonal
-    eye_check = node.inputs[0]
+    potential_eye = 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
+        potential_eye.owner
+        and isinstance(potential_eye.owner.op, Eye)
+        and hasattr(potential_eye.owner.inputs[-1], "data")
+        and potential_eye.owner.inputs[-1].data.item() == 0
     ):
         return None
-    return [eye_check]
+    return [potential_eye]
 
 
 @register_canonicalize
@@ -665,10 +666,9 @@ def rewrite_cholesky_diag_to_sqrt_diag(fgraph, node):
         and isinstance(inputs.owner.op, AllocDiag)
         and AllocDiag.is_offset_zero(inputs.owner)
     ):
-        cholesky_input = inputs.owner.inputs[0]
-        if cholesky_input.type.ndim == 1:
-            cholesky_val = pt.diag(cholesky_input**0.5)
-            return [cholesky_val]
+        diag_input = inputs.owner.inputs[0]
+        cholesky_val = pt.diag(diag_input**0.5)
+        return [cholesky_val]
 
     # Check if the input is an elemwise multiply with identity matrix -- this also results in a diagonal matrix
     inputs_or_none = _find_diag_from_eye_mul(inputs)
@@ -686,8 +686,6 @@ def rewrite_cholesky_diag_to_sqrt_diag(fgraph, node):
     # Now, we can simply return the matrix consisting of sqrt values of the original diagonal elements
     # For a matrix, we have to first extract the diagonal (non-zero values) and then only use those
     if non_eye_input.type.broadcastable[-2:] == (False, False):
-        # For Matrix
-        return [eye_input * (non_eye_input.diagonal(axis1=-1, axis2=-2) ** 0.5)]
-    else:
-        # For Vector or Scalar
-        return [eye_input * (non_eye_input**0.5)]
+        non_eye_input = non_eye_input.diagonal(axis1=-1, axis2=-2)
+
+    return [eye_input * (non_eye_input**0.5)]

tests/tensor/rewriting/test_linalg.py

@@ -572,17 +572,12 @@ def get_pt_function(x, op_name):
 
 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)
@@ -656,3 +651,15 @@ def test_cholesky_diag_from_diag():
         atol=1e-3 if config.floatX == "float32" else 1e-8,
         rtol=1e-3 if config.floatX == "float32" else 1e-8,
     )
+
+
+def test_dont_apply_cholesky():
+    x = pt.tensor("x", shape=(7, 7))
+    y = pt.eye(7, k=-1) * x
+    # Here, y is not a diagonal matrix because of k = -1
+    z_cholesky = pt.linalg.cholesky(y)
+
+    # REWRITE TEST (should not be applied)
+    f_rewritten = function([x], z_cholesky, mode="FAST_RUN")
+    nodes = f_rewritten.maker.fgraph.apply_nodes
+    assert any(isinstance(node.op, Cholesky) for node in nodes)