Simplify linalg rewrites with pattern matching

Author: ricardoV94

pytensor/tensor/_linalg/solve/rewriting.py

@@ -15,7 +15,6 @@
 from pytensor.tensor.elemwise import DimShuffle
 from pytensor.tensor.rewriting.basic import register_specialize
 from pytensor.tensor.rewriting.blockwise import blockwise_of
-from pytensor.tensor.rewriting.linalg import is_matrix_transpose
 from pytensor.tensor.slinalg import Solve, cho_solve, cholesky, lu_factor, lu_solve
 from pytensor.tensor.variable import TensorVariable
 
@@ -79,14 +78,16 @@ def get_root_A(a: TensorVariable) -> tuple[TensorVariable, bool]:
         # the root variable is the pre-DimShuffled input.
         # Otherwise, `a` is considered the root variable.
         # We also return whether the root `a` is transposed.
-        transposed = False
-        if a.owner is not None and isinstance(a.owner.op, DimShuffle):
-            if a.owner.op.is_left_expand_dims:
-                [a] = a.owner.inputs
-            elif is_matrix_transpose(a):
-                [a] = a.owner.inputs
+        match a.owner_op_and_inputs:
+            case (DimShuffle(is_left_expand_dims=True), root_a):
+                transposed = False
+            case (DimShuffle(is_left_expanded_matrix_transpose=True), root_a):
                 transposed = True
-        return a, transposed
+            case _:
+                root_a = a
+                transposed = False
+
+        return root_a, transposed
 
     def find_solve_clients(var, assume_a):
         clients = []
@@ -119,11 +120,11 @@ def find_solve_clients(var, assume_a):
 
     # Find Solves using A.T
     for cl, _ in fgraph.clients[A]:
-        if isinstance(cl.op, DimShuffle) and is_matrix_transpose(cl.out):
-            A_T = cl.out
-            A_solve_clients_and_transpose.extend(
-                (client, True) for client in find_solve_clients(A_T, assume_a)
-            )
+        match (cl.op, *cl.outputs):
+            case (DimShuffle(is_left_expanded_matrix_transpose=True), A_T):
+                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 theres' a single use don't do it... unless it's being broadcast in a Blockwise (or we're eager)