added rewrite for cholesky(diag) -> sqrt(diag)

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -572,3 +572,28 @@ def rewrite_cholesky_eye_to_eye(fgraph, node):
     ):
         return None
     return [eye_check]
+
+
+@register_canonicalize
+@register_stabilize
+@node_rewriter([Blockwise])
+def rewrite_cholesky_diag_from_eye_mul(fgraph, node):
+    # Find whether cholesky op is being applied
+    if not isinstance(node.op.core_op, Cholesky):
+        return None
+
+    # Check whether input is diagonal from multiplcation of identity matrix with a tensor
+    inputs = node.inputs[0]
+    inputs_or_none = _find_diag_from_eye_mul(inputs)
+    if inputs_or_none is None:
+        return None
+
+    eye_input, non_eye_inputs = inputs_or_none
+
+    # Dealing with only one other input
+    if len(non_eye_inputs) != 1:
+        return None
+
+    eye_input, non_eye_input = eye_input[0], non_eye_inputs[0]
+
+    return [eye_input * (non_eye_input**0.5)]

tests/tensor/rewriting/test_linalg.py

@@ -572,3 +572,39 @@ def test_cholesky_eye_rewrite():
         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,), (1, 7), (7, 1), (7, 7), (3, 7, 7)],
+    ids=["scalar", "vector", "row_vec", "col_vec", "matrix", "batched_input"],
+)
+def test_cholesky_diag_from_eye_mul(shape):
+    # Initializing x based on scalar/vector/matrix
+    x = pt.tensor("x", shape=shape)
+    y = pt.eye(7) * x
+    # Performing cholesky decomposition using pt.linalg.cholesky
+    z_cholesky = pt.linalg.cholesky(y)
+
+    # REWRITE TEST
+    f_rewritten = function([x], z_cholesky, mode="FAST_RUN")
+    nodes = f_rewritten.maker.fgraph.apply_nodes
+    assert not any(isinstance(node.op, Cholesky) for node in nodes)
+
+    # NUMERIC VALUE TEST
+    if len(shape) == 0:
+        x_test = np.array(np.random.rand()).astype(config.floatX)
+    elif len(shape) == 1:
+        x_test = np.random.rand(*shape).astype(config.floatX)
+    else:
+        x_test = np.random.rand(*shape).astype(config.floatX)
+    x_test_matrix = np.eye(7) * x_test
+    cholesky_val = np.linalg.cholesky(x_test_matrix)
+    rewritten_val = f_rewritten(x_test)
+
+    assert_allclose(
+        cholesky_val,
+        rewritten_val,
+        atol=1e-3 if config.floatX == "float32" else 1e-8,
+        rtol=1e-3 if config.floatX == "float32" else 1e-8,
+    )