Tweak transpose

Author: ricardoV94

pytensor/xtensor/shape.py

@@ -6,7 +6,7 @@
 from pytensor.graph import Apply
 from pytensor.scalar import upcast
 from pytensor.xtensor.basic import XOp
-from pytensor.xtensor.type import as_xtensor, xtensor
+from pytensor.xtensor.type import XTensorVariable, as_xtensor, xtensor
 
 
 class Stack(XOp):
@@ -75,102 +75,55 @@ 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")
+    __props__ = ("dims",)
 
     def __init__(
         self,
         dims: tuple[str | Literal[...], ...],
-        missing_dims: Literal["raise", "warn", "ignore"] = "raise",
     ):
         super().__init__()
+        if dims.count(...) > 1:
+            raise ValueError("an index can only have a single ellipsis ('...')")
         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
-        )
+
+        transpose_dims = self.dims
+        x_dims = x.type.dims
+
+        if transpose_dims == () or transpose_dims == (...,):
+            out_dims = tuple(reversed(x_dims))
+        elif ... in transpose_dims:
+            # Handle ellipsis expansion
+            ellipsis_idx = transpose_dims.index(...)
+            pre = transpose_dims[:ellipsis_idx]
+            post = transpose_dims[ellipsis_idx + 1 :]
+            middle = [d for d in x_dims if d not in pre + post]
+            out_dims = (*pre, *middle, *post)
+            if set(out_dims) != set(x_dims):
+                raise ValueError(f"{out_dims} must be a permuted list of {x_dims}")
+        else:
+            out_dims = transpose_dims
+            if set(out_dims) != set(x_dims):
+                raise ValueError(
+                    f"{out_dims} must be a permuted list of {x_dims}, unless `...` is included"
+                )
 
         output = xtensor(
             dtype=x.type.dtype,
-            shape=tuple(x.type.shape[x.type.dims.index(d)] for d in dims),
-            dims=dims,
+            shape=tuple(x.type.shape[x.type.dims.index(d)] for d in out_dims),
+            dims=out_dims,
         )
         return Apply(self, [x], [output])
 
 
-def transpose(x, *dims, missing_dims: Literal["raise", "warn", "ignore"] = "raise"):
+def transpose(
+    x,
+    *dims: str | Literal[...],
+    missing_dims: Literal["raise", "warn", "ignore"] = "raise",
+) -> XTensorVariable:
     """Transpose dimensions of the tensor.
 
     Parameters
@@ -196,7 +149,21 @@ def transpose(x, *dims, missing_dims: Literal["raise", "warn", "ignore"] = "rais
     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)
+    # Validate dimensions
+    x = as_xtensor(x)
+    all_dims = x.type.dims
+    invalid_dims = set(dims) - {..., *all_dims}
+    if invalid_dims:
+        if missing_dims != "ignore":
+            msg = f"Dimensions {invalid_dims} do not exist. Expected one or more of: {all_dims}"
+            if missing_dims == "raise":
+                raise ValueError(msg)
+            else:
+                warnings.warn(msg)
+        # Handle missing dimensions if not raising
+        dims = tuple(d for d in dims if d in all_dims or d is ...)
+
+    return Transpose(dims)(x)
 
 
 class Concat(XOp):

pytensor/xtensor/type.py

@@ -10,7 +10,7 @@
     XARRAY_AVAILABLE = False
 
 from collections.abc import Sequence
-from typing import TypeVar, Union
+from typing import Literal, TypeVar
 
 import numpy as np
 
@@ -357,7 +357,11 @@ 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"):
+    def transpose(
+        self,
+        *dims: str | Literal[...],
+        missing_dims: Literal["raise", "warn", "ignore"] = "raise",
+    ) -> "XTensorVariable":
         """Transpose dimensions of the tensor.
 
         Parameters

tests/xtensor/test_shape.py

@@ -1,4 +1,6 @@
 # ruff: noqa: E402
+import re
+
 import pytest
 
 
@@ -51,6 +53,62 @@ def test_transpose():
         xr_assert_allclose(res_i, expected_res_i)
 
 
+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=re.escape(
+            "Dimensions {'d'} do not exist. Expected one or more of: ('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.T)
+
+
 def test_stack():
     dims = ("a", "b", "c", "d")
     x = xtensor("x", dims=dims, shape=(2, 3, 5, 7))
@@ -153,65 +211,3 @@ 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())