Add transpose operation for labeled tensors with ellipsis support

Author: AllenDowney

pytensor/xtensor/rewriting/shape.py

@@ -2,7 +2,7 @@
 from pytensor.tensor import broadcast_to, join, moveaxis
 from pytensor.xtensor.basic import tensor_from_xtensor, xtensor_from_tensor
 from pytensor.xtensor.rewriting.basic import register_xcanonicalize
-from pytensor.xtensor.shape import Concat, Stack
+from pytensor.xtensor.shape import Concat, Stack, Transpose
 
 
 @register_xcanonicalize
@@ -70,3 +70,35 @@ def lower_concat(fgraph, node):
     joined_tensor = join(concat_axis, *bcast_tensor_inputs)
     new_out = xtensor_from_tensor(joined_tensor, dims=out_dims)
     return [new_out]
+
+
+@register_xcanonicalize
+@node_rewriter(tracks=[Transpose])
+def lower_transpose(fgraph, node):
+    [x] = node.inputs
+    # Determine the permutation of axes
+    out_dims = node.op.dims
+    in_dims = x.type.dims
+    # Expand ellipsis if present
+    if out_dims == () or out_dims == (...,):
+        expanded_dims = tuple(reversed(in_dims))
+    elif ... in out_dims:
+        pre = []
+        post = []
+        found = False
+        for d in out_dims:
+            if d is ...:
+                found = True
+            elif not found:
+                pre.append(d)
+            else:
+                post.append(d)
+        middle = [d for d in in_dims if d not in pre + post]
+        expanded_dims = tuple(pre + middle + post)
+    else:
+        expanded_dims = out_dims
+    perm = tuple(in_dims.index(d) for d in expanded_dims)
+    x_tensor = tensor_from_xtensor(x)
+    x_tensor_transposed = x_tensor.transpose(perm)
+    new_out = xtensor_from_tensor(x_tensor_transposed, dims=expanded_dims)
+    return [new_out]

pytensor/xtensor/shape.py

@@ -73,6 +73,49 @@ def stack(x, dim: dict[str, Sequence[str]] | None = None, **dims: Sequence[str])
     return y
 
 
+class Transpose(XOp):
+    __props__ = ("dims",)
+
+    def __init__(self, dims: tuple[str, ...]):
+        super().__init__()
+        self.dims = dims
+
+    def make_node(self, x):
+        x = as_xtensor(x)
+        # Allow ellipsis for full transpose
+        if self.dims == () or self.dims == (...,):
+            dims = tuple(reversed(x.type.dims))
+        else:
+            # Expand ellipsis if present
+            if ... in self.dims:
+                pre = []
+                post = []
+                found = False
+                for d in self.dims:
+                    if d is ...:
+                        found = True
+                    elif not found:
+                        pre.append(d)
+                    else:
+                        post.append(d)
+                middle = [d for d in x.type.dims if d not in pre + post]
+                dims = tuple(pre + middle + post)
+            else:
+                dims = self.dims
+        if set(dims) != set(x.type.dims):
+            raise ValueError(f"Transpose dims {dims} must match {x.type.dims}")
+        output = xtensor(
+            dtype=x.type.dtype,
+            shape=tuple(x.type.shape[x.type.dims.index(d)] for d in dims),
+            dims=dims,
+        )
+        return Apply(self, [x], [output])
+
+
+def transpose(x, *dims):
+    return Transpose(dims)(x)
+
+
 class Concat(XOp):
     __props__ = ("dim",)
 

tests/xtensor/test_shape.py

@@ -10,7 +10,7 @@
 from xarray import DataArray
 from xarray import concat as xr_concat
 
-from pytensor.xtensor.shape import concat, stack
+from pytensor.xtensor.shape import concat, stack, transpose
 from pytensor.xtensor.type import xtensor
 from tests.xtensor.util import xr_assert_allclose, xr_function, xr_random_like
 
@@ -24,9 +24,8 @@ def powerset(iterable, min_group_size=0):
     )
 
 
-@pytest.mark.xfail(reason="Not yet implemented")
+# @pytest.mark.xfail(reason="Not yet implemented")
 def test_transpose():
-    transpose = None
     a, b, c, d, e = "abcde"
 
     x = xtensor("x", dims=(a, b, c, d, e), shape=(2, 3, 5, 7, 11))