Implement transpose for XTensorVariables

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,19 @@ 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
+    # Use the final dimensions that were already computed in make_node
+    out_dims = node.outputs[0].type.dims
+    in_dims = x.type.dims
+
+    # Compute the permutation based on the final dimensions
+    perm = tuple(in_dims.index(d) for d in out_dims)
+    x_tensor = tensor_from_xtensor(x)
+    x_tensor_transposed = x_tensor.transpose(perm)
+    new_out = xtensor_from_tensor(x_tensor_transposed, dims=out_dims)
+    return [new_out]

pytensor/xtensor/shape.py

@@ -1,4 +1,6 @@
+import warnings
 from collections.abc import Sequence
+from typing import Literal
 
 from pytensor import Variable
 from pytensor.graph import Apply
@@ -73,6 +75,130 @@ def stack(x, dim: dict[str, Sequence[str]] | None = None, **dims: Sequence[str])
     return y
 
 
+def expand_ellipsis(
+    dims: tuple[str, ...],
+    all_dims: tuple[str, ...],
+    validate: bool = True,
+    missing_dims: Literal["raise", "warn", "ignore"] = "raise",
+) -> tuple[str, ...]:
+    """Expand ellipsis in dimension permutation.
+
+    Parameters
+    ----------
+    dims : tuple[str, ...]
+        The dimension permutation, which may contain ellipsis
+    all_dims : tuple[str, ...]
+        All available dimensions
+    validate : bool, default True
+        Whether to check that all non-ellipsis elements in dims are valid dimension names.
+    missing_dims : {"raise", "warn", "ignore"}, optional
+        How to handle dimensions that don't exist in all_dims:
+        - "raise": Raise an error if any dimensions don't exist (default)
+        - "warn": Warn if any dimensions don't exist
+        - "ignore": Silently ignore any dimensions that don't exist
+
+    Returns
+    -------
+    tuple[str, ...]
+        The expanded dimension permutation
+
+    Raises
+    ------
+    ValueError
+        If more than one ellipsis is present in dims.
+        If any non-ellipsis element in dims is not a valid dimension name and validate is True.
+        If missing_dims is "raise" and any dimension in dims doesn't exist in all_dims.
+    """
+    # Handle empty or full ellipsis case
+    if dims == () or dims == (...,):
+        return tuple(reversed(all_dims))
+
+    # Check for multiple ellipses
+    if dims.count(...) > 1:
+        raise ValueError("an index can only have a single ellipsis ('...')")
+
+    # Validate dimensions if requested
+    if validate:
+        invalid_dims = set(dims) - {..., *all_dims}
+        if invalid_dims:
+            if missing_dims == "raise":
+                raise ValueError(
+                    f"Invalid dimensions: {invalid_dims}. Available dimensions: {all_dims}"
+                )
+            elif missing_dims == "warn":
+                warnings.warn(f"Dimensions {invalid_dims} do not exist in {all_dims}")
+
+    # Handle missing dimensions if not raising
+    if missing_dims in ("ignore", "warn"):
+        dims = tuple(d for d in dims if d in all_dims or d is ...)
+
+    # If no ellipsis, just return the dimensions
+    if ... not in dims:
+        return dims
+
+    # Handle ellipsis expansion
+    ellipsis_idx = dims.index(...)
+    pre = list(dims[:ellipsis_idx])
+    post = list(dims[ellipsis_idx + 1 :])
+    middle = [d for d in all_dims if d not in pre + post]
+    return tuple(pre + middle + post)
+
+
+class Transpose(XOp):
+    __props__ = ("dims", "missing_dims")
+
+    def __init__(
+        self,
+        dims: tuple[str | Literal[...], ...],
+        missing_dims: Literal["raise", "warn", "ignore"] = "raise",
+    ):
+        super().__init__()
+        self.dims = dims
+        self.missing_dims = missing_dims
+
+    def make_node(self, x):
+        x = as_xtensor(x)
+        dims = expand_ellipsis(
+            self.dims, x.type.dims, validate=True, missing_dims=self.missing_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, missing_dims: Literal["raise", "warn", "ignore"] = "raise"):
+    """Transpose dimensions of the tensor.
+
+    Parameters
+    ----------
+    x : XTensorVariable
+        Input tensor to transpose.
+    *dims : str
+        Dimensions to transpose to. Can include ellipsis (...) to represent
+        remaining dimensions in their original order.
+    missing_dims : {"raise", "warn", "ignore"}, optional
+        How to handle dimensions that don't exist in the input tensor:
+        - "raise": Raise an error if any dimensions don't exist (default)
+        - "warn": Warn if any dimensions don't exist
+        - "ignore": Silently ignore any dimensions that don't exist
+
+    Returns
+    -------
+    XTensorVariable
+        Transposed tensor with reordered dimensions.
+
+    Raises
+    ------
+    ValueError
+        If any dimension in dims doesn't exist in the input tensor and missing_dims is "raise".
+    """
+    return Transpose(dims, missing_dims=missing_dims)(x)
+
+
 class Concat(XOp):
     __props__ = ("dim",)
 

pytensor/xtensor/type.py

@@ -10,7 +10,7 @@
     XARRAY_AVAILABLE = False
 
 from collections.abc import Sequence
-from typing import TypeVar
+from typing import TypeVar, Union
 
 import numpy as np
 
@@ -357,6 +357,46 @@ def imag(self):
     def real(self):
         return px.math.real(self)
 
+    def transpose(self, *dims: Union[str, type(Ellipsis)], missing_dims: Literal["raise", "warn", "ignore"] = "raise"):
+        """Transpose dimensions of the tensor.
+
+        Parameters
+        ----------
+        *dims : str | Ellipsis
+            Dimensions to transpose. If empty, performs a full transpose.
+            Can use ellipsis (...) to represent remaining dimensions.
+        missing_dims : {"raise", "warn", "ignore"}, default="raise"
+            How to handle dimensions that don't exist in the tensor:
+            - "raise": Raise an error if any dimensions don't exist
+            - "warn": Warn if any dimensions don't exist
+            - "ignore": Silently ignore any dimensions that don't exist
+
+        Returns
+        -------
+        XTensorVariable
+            Transposed tensor with reordered dimensions.
+
+        Raises
+        ------
+        ValueError
+            If missing_dims="raise" and any dimensions don't exist.
+            If multiple ellipsis are provided.
+        """
+        return px.shape.transpose(self, *dims, missing_dims=missing_dims)
+
+    @property
+    def T(self) -> "XTensorVariable":
+        """Return the full transpose of the tensor.
+
+        This is equivalent to calling transpose() with no arguments.
+
+        Returns
+        -------
+        XTensorVariable
+            Fully transposed tensor.
+        """
+        return self.transpose()
+
     # Aggregation
     # https://docs.xarray.dev/en/latest/api.html#id6
     def all(self, dim):

tests/xtensor/test_shape.py

@@ -9,7 +9,7 @@
 import numpy as np
 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_arange_like,
@@ -28,9 +28,7 @@ def powerset(iterable, min_group_size=0):
     )
 
 
-@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))
@@ -155,3 +153,65 @@ def test_concat_scalar():
     res = fn(x1_test, x2_test)
     expected_res = xr_concat([x1_test, x2_test], dim="new_dim")
     xr_assert_allclose(res, expected_res)
+
+
+def test_xtensor_variable_transpose():
+    """Test the transpose() method of XTensorVariable."""
+    x = xtensor("x", dims=("a", "b", "c"), shape=(2, 3, 4))
+
+    # Test basic transpose
+    out = x.transpose()
+    fn = xr_function([x], out)
+    x_test = xr_arange_like(x)
+    xr_assert_allclose(fn(x_test), x_test.transpose())
+
+    # Test transpose with specific dimensions
+    out = x.transpose("c", "a", "b")
+    fn = xr_function([x], out)
+    xr_assert_allclose(fn(x_test), x_test.transpose("c", "a", "b"))
+
+    # Test transpose with ellipsis
+    out = x.transpose("c", ...)
+    fn = xr_function([x], out)
+    xr_assert_allclose(fn(x_test), x_test.transpose("c", ...))
+
+    # Test error cases
+    with pytest.raises(
+        ValueError,
+        match="Invalid dimensions: {'d'}. Available dimensions: \\('a', 'b', 'c'\\)",
+    ):
+        x.transpose("d")
+
+    with pytest.raises(ValueError, match="an index can only have a single ellipsis"):
+        x.transpose("a", ..., "b", ...)
+
+    # Test missing_dims parameter
+    # Test ignore
+    out = x.transpose("c", ..., "d", missing_dims="ignore")
+    fn = xr_function([x], out)
+    xr_assert_allclose(fn(x_test), x_test.transpose("c", ...))
+
+    # Test warn
+    with pytest.warns(UserWarning, match="Dimensions {'d'} do not exist"):
+        out = x.transpose("c", ..., "d", missing_dims="warn")
+        fn = xr_function([x], out)
+        xr_assert_allclose(fn(x_test), x_test.transpose("c", ...))
+
+
+def test_xtensor_variable_T():
+    """Test the T property of XTensorVariable."""
+    # Test T property with 3D tensor
+    x = xtensor("x", dims=("a", "b", "c"), shape=(2, 3, 4))
+    out = x.T
+
+    fn = xr_function([x], out)
+    x_test = xr_arange_like(x)
+    xr_assert_allclose(fn(x_test), x_test.transpose())
+
+    # Test T property with 2D tensor
+    x = xtensor("x", dims=("a", "b"), shape=(2, 3))
+    out = x.T
+
+    fn = xr_function([x], out)
+    x_test = xr_arange_like(x)
+    xr_assert_allclose(fn(x_test), x_test.transpose())