Handle lower flags more carefully

Author: jessegrabowski

pytensor/tensor/_linalg/solve/rewriting.py

@@ -100,40 +100,43 @@ def find_solve_clients(var, assume_a):
             elif isinstance(cl.op, DimShuffle) and cl.op.is_left_expand_dims:
                 # If it's a left expand_dims, recurse on the output
                 clients.extend(find_solve_clients(cl.outputs[0], assume_a))
+
         return clients
 
     assume_a = node.op.core_op.assume_a
 
     if assume_a not in allowed_assume_a:
         return None
 
-    A, _ = get_root_A(node.inputs[0])
+    root_A, root_A_transposed = get_root_A(node.inputs[0])
 
     # Find Solve using A (or left expand_dims of A)
     # TODO: We could handle arbitrary shuffle of the batch dimensions, just need to propagate
     #  that to the A_decomp outputs
-    A_solve_clients_and_transpose = [
-        (client, False) for client in find_solve_clients(A, assume_a)
+    root_A_solve_clients_and_transpose = [
+        (client, False) for client in find_solve_clients(root_A, assume_a)
     ]
 
     # Find Solves using A.T
-    for cl, _ in fgraph.clients[A]:
+    for cl, _ in fgraph.clients[root_A]:
         if isinstance(cl.op, DimShuffle) and is_matrix_transpose(cl.out):
             A_T = cl.out
-            A_solve_clients_and_transpose.extend(
+            root_A_solve_clients_and_transpose.extend(
                 (client, True) for client in find_solve_clients(A_T, assume_a)
             )
 
-    if not eager and len(A_solve_clients_and_transpose) == 1:
+    if not eager and len(root_A_solve_clients_and_transpose) == 1:
         # If theres' a single use don't do it... unless it's being broadcast in a Blockwise (or we're eager)
         # That's a "reuse" inside the inner vectorized loop
         batch_ndim = node.op.batch_ndim(node)
-        (client, _) = A_solve_clients_and_transpose[0]
-        original_A, b = client.inputs
+        (client, _) = root_A_solve_clients_and_transpose[0]
+
+        A, b = client.inputs
+
         if not any(
             a_bcast and not b_bcast
             for a_bcast, b_bcast in zip(
-                original_A.type.broadcastable[:batch_ndim],
+                A.type.broadcastable[:batch_ndim],
                 b.type.broadcastable[:batch_ndim],
                 strict=True,
             )
@@ -142,19 +145,27 @@ def find_solve_clients(var, assume_a):
 
     # If any Op had check_finite=True, we also do it for the LU decomposition
     check_finite_decomp = False
-    for client, _ in A_solve_clients_and_transpose:
+    for client, _ in root_A_solve_clients_and_transpose:
         if client.op.core_op.check_finite:
             check_finite_decomp = True
             break
 
-    lower = node.op.core_op.lower
+    (first_solve, transposed) = root_A_solve_clients_and_transpose[0]
+    lower = first_solve.op.core_op.lower
+    if transposed:
+        lower = not lower
+
     A_decomp = decompose_A(
-        A, assume_a=assume_a, check_finite=check_finite_decomp, lower=lower
+        root_A, assume_a=assume_a, check_finite=check_finite_decomp, lower=lower
     )
 
     replacements = {}
-    for client, transposed in A_solve_clients_and_transpose:
+    for client, transposed in root_A_solve_clients_and_transpose:
         _, b = client.inputs
+        lower = client.op.core_op.lower
+        if transposed:
+            lower = not lower
+
         new_x = solve_decomposed_system(
             A_decomp,
             b,

pytensor/tensor/slinalg.py

@@ -201,7 +201,9 @@ def cholesky(
 
     """
 
-    return Blockwise(Cholesky(lower=lower, on_error=on_error))(x)
+    return Blockwise(
+        Cholesky(lower=lower, on_error=on_error, check_finite=check_finite)
+    )(x)
 
 
 class SolveBase(Op):

tests/tensor/linalg/test_rewriting.py

@@ -112,14 +112,15 @@ def test_lu_decomposition_reused_forward_and_gradient(assume_a, counter, transpo
     np.testing.assert_allclose(resg0, resg1, rtol=rtol)
 
 
-@pytest.mark.parametrize("transposed", (False, True))
+@pytest.mark.parametrize("transposed", (False, True), ids=["no_trans", "trans"])
 @pytest.mark.parametrize(
     "assume_a, counter",
     (
         ("gen", LUOpCounter),
         ("tridiagonal", TriDiagLUOpCounter),
         ("pos", CholeskyOpCounter),
     ),
+    ids=["assume_gen", "assume_tridiagonal", "assume_pos"],
 )
 def test_lu_decomposition_reused_blockwise(assume_a, counter, transposed):
     rewrite_name = reuse_decomposition_multiple_solves.__name__
@@ -251,3 +252,39 @@ def test_decomposition_reused_preserves_check_finite(assume_a, counter):
         assert fn_opt(A_valid, b1_valid * np.nan, b2_valid)
     with pytest.raises(ValueError, match="array must not contain infs or NaNs"):
         assert fn_opt(A_valid * np.nan, b1_valid, b2_valid)
+
+
+@pytest.mark.parametrize(
+    "lower_first", [True, False], ids=["lower_first", "upper_first"]
+)
+def test_cho_solve_handles_lower_flags(lower_first):
+    A = tensor("A", shape=(2, None))
+    b = tensor("b", shape=(2,))
+
+    x1 = solve(A, b, assume_a="pos", lower=lower_first, check_finite=False)
+    x2 = solve(A.mT, b, assume_a="pos", lower=not lower_first, check_finite=False)
+
+    dx1_dA = grad(x1.sum(), A)
+    dx2_dA = grad(x2.sum(), A)
+
+    fn = function([A, b], [x1, dx1_dA, x2, dx2_dA])
+
+    rng = np.random.default_rng()
+    L_values = rng.normal(size=(2, 2))
+    A_values = L_values @ L_values.T  # Ensure A is positive definite
+
+    if lower_first:
+        A_values[0, 1] = np.nan
+    else:
+        A_values[1, 0] = np.nan
+
+    b_values = rng.normal(size=(2,))
+
+    # This computation should not raise an error, and none of them should be NaN
+    res = fn(A_values, b_values)
+    for x in res:
+        assert np.isfinite(x).all()
+
+    # If we put the NaN in the wrong place, it should raise an error
+    with pytest.raises(np.linalg.LinAlgError):
+        fn(A_values.T, b_values)