Cache is_matrix_transpose at DimShuffle Op level

Author: ricardoV94

pytensor/tensor/elemwise.py

@@ -159,8 +159,7 @@ def __init__(self, *, input_ndim: int, new_order: Sequence[int | Literal["x"]]):
         # This is the list of the original dimensions that we keep
         self.shuffle = [x for x in new_order if x != "x"]
         self.transposition = self.shuffle + drop
-        # List of dimensions of the output that are broadcastable and were not
-        # in the original input
+        # List of dimensions of the output that are broadcastable and were not in the original input
         self.augment = augment = sorted(i for i, x in enumerate(new_order) if x == "x")
         self.drop = drop
 
@@ -175,6 +174,12 @@ def __init__(self, *, input_ndim: int, new_order: Sequence[int | Literal["x"]]):
         self.is_right_expand_dims = self.is_expand_dims and new_order[
             :input_ndim
         ] == list(range(input_ndim))
+        self.is_matrix_transpose = False
+        if dims_are_shuffled and (not drop) and input_ndim >= 2:
+            # We consider a matrix transpose if we only flip the last two dims
+            # Regardless of whethre there's an expand_dims or not
+            mt_pattern = [*range(input_ndim - 2), input_ndim - 1, input_ndim - 2]
+            self.is_matrix_transpose = new_order[len(augment) :] == mt_pattern
 
     def __setstate__(self, state):
         self.__dict__.update(state)

pytensor/tensor/rewriting/linalg.py

@@ -66,25 +66,11 @@
 def is_matrix_transpose(x: TensorVariable) -> bool:
     """Check if a variable corresponds to a transpose of the last two axes"""
     node = x.owner
-    if (
-        node
+    return (
+        node is not None
         and isinstance(node.op, DimShuffle)
-        and not (node.op.drop or node.op.augment)
-    ):
-        [inp] = node.inputs
-        ndims = inp.type.ndim
-        if ndims < 2:
-            return False
-        transpose_order = (*range(ndims - 2), ndims - 1, ndims - 2)
-
-        # Allow expand_dims on the left of the transpose
-        if (diff := len(transpose_order) - len(node.op.new_order)) > 0:
-            transpose_order = (
-                *(["x"] * diff),
-                *transpose_order,
-            )
-        return node.op.new_order == transpose_order
-    return False
+        and node.op.is_matrix_transpose
+    )
 
 
 @register_canonicalize