More carefully handle lower flag in Solve

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,