added test for batched case and more cases of not applying rewrite

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -962,5 +962,7 @@ def rewrite_cholesky_diag_to_sqrt_diag(fgraph, node):
     # 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):
         non_eye_input = non_eye_input.diagonal(axis1=-1, axis2=-2)
+        if eye_input.type.ndim > 2:
+            non_eye_input = pt.shape_padaxis(non_eye_input, -2)
 
     return [eye_input * (non_eye_input**0.5)]

tests/tensor/rewriting/test_linalg.py

@@ -829,8 +829,8 @@ def test_cholesky_eye_rewrite():
 
 @pytest.mark.parametrize(
     "shape",
-    [(), (7,), (7, 7)],
-    ids=["scalar", "vector", "matrix"],
+    [(), (7,), (7, 7), (5, 7, 7)],
+    ids=["scalar", "vector", "matrix", "batched"],
 )
 def test_cholesky_diag_from_eye_mul(shape):
     # Initializing x based on scalar/vector/matrix
@@ -888,13 +888,21 @@ def test_cholesky_diag_from_diag():
     )
 
 
-def test_dont_apply_cholesky():
+def test_rewrite_cholesky_diag_to_sqrt_diag_not_applied():
+    # Case 1 : y is not a diagonal matrix because of k = -1
     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)
+
+    # Case 2 : eye is degenerate
+    x = pt.scalar("x")
+    y = pt.eye(1) * x
+    z_cholesky = pt.linalg.cholesky(y)
+    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)