Simplify linalg rewrites with pattern matching

Author: ricardoV94

pytensor/tensor/_linalg/solve/rewriting.py

@@ -92,16 +92,13 @@ def get_root_A(a: TensorVariable) -> tuple[TensorVariable, bool]:
     def find_solve_clients(var, assume_a):
         clients = []
         for cl, idx in fgraph.clients[var]:
-            if (
-                idx == 0
-                and isinstance(cl.op, Blockwise)
-                and isinstance(cl.op.core_op, Solve)
-                and (cl.op.core_op.assume_a == assume_a)
-            ):
-                clients.append(cl)
-            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))
+            match (idx, cl.op, *cl.outputs):
+                case (0, Blockwise(Solve(assume_a=assume_a_var)), *_) if (
+                    assume_a_var == assume_a
+                ):
+                    clients.append(cl)
+                case (0, DimShuffle(is_left_expand_dims=True), cl_out):
+                    clients.extend(find_solve_clients(cl_out, assume_a))
         return clients
 
     assume_a = node.op.core_op.assume_a
@@ -186,34 +183,34 @@ def _scan_split_non_sequence_decomposition_and_solve(
     changed = False
     while True:
         for inner_node in new_scan_fgraph.toposort():
-            if (
-                isinstance(inner_node.op, Blockwise)
-                and isinstance(inner_node.op.core_op, Solve)
-                and inner_node.op.core_op.assume_a in allowed_assume_a
-            ):
-                A, _b = inner_node.inputs
-                if all(
-                    (isinstance(root_inp, Constant) or (root_inp in non_sequences))
-                    for root_inp in graph_inputs([A])
+            match (inner_node.op, *inner_node.inputs):
+                case (Blockwise(Solve(assume_a=assume_a_var)), A, _b) if (
+                    assume_a_var in allowed_assume_a
                 ):
-                    if new_scan_fgraph is scan_op.fgraph:
-                        # Clone the first time to avoid mutating the original fgraph
-                        new_scan_fgraph, equiv = new_scan_fgraph.clone_get_equiv()
-                        non_sequences = {equiv[non_seq] for non_seq in non_sequences}
-                        inner_node = equiv[inner_node]  # type: ignore
-
-                    replace_dict = _split_decomp_and_solve_steps(
-                        new_scan_fgraph,
-                        inner_node,
-                        eager=True,
-                        allowed_assume_a=allowed_assume_a,
-                    )
-                    assert isinstance(replace_dict, dict) and len(replace_dict) > 0, (
-                        "Rewrite failed"
-                    )
-                    new_scan_fgraph.replace_all(replace_dict.items())
-                    changed = True
-                    break  # Break to start over with a fresh toposort
+                    if all(
+                        (isinstance(root_inp, Constant) or (root_inp in non_sequences))
+                        for root_inp in graph_inputs([A])
+                    ):
+                        if new_scan_fgraph is scan_op.fgraph:
+                            # Clone the first time to avoid mutating the original fgraph
+                            new_scan_fgraph, equiv = new_scan_fgraph.clone_get_equiv()
+                            non_sequences = {
+                                equiv[non_seq] for non_seq in non_sequences
+                            }
+                            inner_node = equiv[inner_node]  # type: ignore
+
+                        replace_dict = _split_decomp_and_solve_steps(
+                            new_scan_fgraph,
+                            inner_node,
+                            eager=True,
+                            allowed_assume_a=allowed_assume_a,
+                        )
+                        assert (
+                            isinstance(replace_dict, dict) and len(replace_dict) > 0
+                        ), "Rewrite failed"
+                        new_scan_fgraph.replace_all(replace_dict.items())
+                        changed = True
+                        break  # Break to start over with a fresh toposort
         else:  # no_break
             break  # Nothing else changed
 

pytensor/tensor/rewriting/linalg.py

@@ -2,7 +2,6 @@
 
 import numpy as np
 
-from pytensor import Variable
 from pytensor import tensor as pt
 from pytensor.compile import optdb
 from pytensor.graph import Apply, Constant, FunctionGraph
@@ -310,9 +309,7 @@ def log_of_prod_to_sum_of_log(fgraph, node):
 
 @register_specialize
 @node_rewriter([blockwise_of(MatrixInverse | Cholesky | MatrixPinv)])
-def lift_linalg_of_expanded_matrices(
-    fgraph: FunctionGraph, node: Apply
-) -> list[Variable] | None:
+def lift_linalg_of_expanded_matrices(fgraph: FunctionGraph, node: Apply):
     """
     Rewrite compositions of linear algebra operations by lifting expensive operations (Cholesky, Inverse) through Ops
     that join matrices (KroneckerProduct, BlockDiagonal).

tests/tensor/linalg/test_rewriting.py

@@ -86,13 +86,21 @@ def test_lu_decomposition_reused_forward_and_gradient(assume_a, counter, transpo
 
     x = solve(A, b, assume_a=assume_a, transposed=transposed)
     grad_x_wrt_A = grad(x.sum(), A)
-    fn_no_opt = function([A, b], [x, grad_x_wrt_A], mode=mode.excluding(rewrite_name))
+    fn_no_opt = function(
+        [A, b],
+        [x, grad_x_wrt_A],
+        mode=mode.excluding(rewrite_name, "psd_solve_to_chol_solve"),
+    )
     no_opt_nodes = fn_no_opt.maker.fgraph.apply_nodes
     assert counter.count_vanilla_solve_nodes(no_opt_nodes) == 2
     assert counter.count_decomp_nodes(no_opt_nodes) == 0
     assert counter.count_solve_nodes(no_opt_nodes) == 0
 
-    fn_opt = function([A, b], [x, grad_x_wrt_A], mode=mode.including(rewrite_name))
+    fn_opt = function(
+        [A, b],
+        [x, grad_x_wrt_A],
+        mode=mode.including(rewrite_name).excluding("psd_solve_to_chol_solve"),
+    )
     opt_nodes = fn_opt.maker.fgraph.apply_nodes
     assert counter.count_vanilla_solve_nodes(opt_nodes) == 0
     assert counter.count_decomp_nodes(opt_nodes) == 1
@@ -129,13 +137,19 @@ def test_lu_decomposition_reused_blockwise(assume_a, counter, transposed):
     b = tensor("b", shape=(2, 3, 4))
 
     x = solve(A, b, assume_a=assume_a, transposed=transposed)
-    fn_no_opt = function([A, b], [x], mode=mode.excluding(rewrite_name))
+    fn_no_opt = function(
+        [A, b], [x], mode=mode.excluding(rewrite_name, "psd_solve_to_chol_solve")
+    )
     no_opt_nodes = fn_no_opt.maker.fgraph.apply_nodes
     assert counter.count_vanilla_solve_nodes(no_opt_nodes) == 1
     assert counter.count_decomp_nodes(no_opt_nodes) == 0
     assert counter.count_solve_nodes(no_opt_nodes) == 0
 
-    fn_opt = function([A, b], [x], mode=mode.including(rewrite_name))
+    fn_opt = function(
+        [A, b],
+        [x],
+        mode=mode.including(rewrite_name).excluding("psd_solve_to_chol_solve"),
+    )
     opt_nodes = fn_opt.maker.fgraph.apply_nodes
     assert counter.count_vanilla_solve_nodes(opt_nodes) == 0
     assert counter.count_decomp_nodes(opt_nodes) == 1
@@ -176,6 +190,7 @@ def test_lu_decomposition_reused_scan(assume_a, counter, transposed):
         non_sequences=[A],
         n_steps=10,
         return_updates=False,
+        mode=get_default_mode().excluding("psd_solve_to_chol_solve"),
     )
 
     fn_no_opt = function(