Add transpose() for labeled tensors #1427

AllenDowney · 2025-05-27T19:01:18Z

Add transpose operation for labeled tensors

This PR implements the transpose operation for labeled tensors (XTensor), which was previously marked as not implemented in the test suite. The implementation follows the xarray-like API for labeled tensors and includes support for ellipsis in dimension permutations.

Changes

Added Transpose XOp class in pytensor/xtensor/shape.py that handles dimension reordering while preserving labels
Added lower_transpose rewrite rule in pytensor/xtensor/rewriting/shape.py to convert labeled tensor operations into regular tensor operations
Added expand_ellipsis helper function to handle ellipsis expansion in dimension permutations
Updated test_transpose in tests/xtensor/test_shape.py to verify the implementation

Features

Support for full transpose (reversing all dimensions)
Support for partial permutations
Support for ellipsis (...) in dimension permutations
Proper handling of dimension labels during transposition
Conversion to regular tensor operations for evaluation

Testing

The implementation passes all test cases, including:

Identity permutation
Full transpose
Empty permutation (equivalent to full transpose)
Swapping last two dimensions
Using ellipsis for partial permutations

xarray has a missing_dims, that we could provide here as well: https://docs.xarray.dev/en/latest/generated/xarray.DataArray.transpose.html

We already do that for isel:

pytensor/pytensor/xtensor/type.py

Lines 362 to 368 in e32d865

def isel(

self,

indexers: dict[str, Any] | None = None,

drop: bool = False, # Unused by PyTensor

missing_dims: Literal["raise", "warn", "ignore"] = "raise",

**indexers_kwargs,

):

ricardoV94 · 2025-05-27T22:29:12Z

pytensor/xtensor/shape.py

@@ -73,6 +73,67 @@ def stack(x, dim: dict[str, Sequence[str]] | None = None, **dims: Sequence[str])
    return y


+def expand_ellipsis(


Should this check that there's at most one dims and raise otherwise?

… matching xarray behavior; update tests and rewrites accordingly

…ent formatting

AllenDowney · 2025-05-28T15:07:43Z

@ricardoV94 I think I've addressed all of your comments -- please take another look.

ricardoV94

Looks pretty good. Just some issue with the type hint, and because of that I took the liberty of leaving some other nitpicky comments. Feel free to not address them if you disagree / it's too cumbersome

pytensor/xtensor/shape.py

tests/xtensor/test_shape.py

… tests

…yle compliance

AllenDowney · 2025-05-28T19:27:10Z

This PR adds support for expanding dimensions in labeled tensors, similar to xarray's functionality. The key changes include:

Added ExpandDims operation that allows adding new dimensions to a tensor while preserving the existing labeled dimensions
Implemented lower_expand_dims rewrite that efficiently handles the dimension expansion using reshape operations
Added corresponding tests to verify the behavior matches xarray's expand_dims functionality
Ensured code style compliance and improved documentation

The implementation allows users to add new dimensions to their labeled tensors while maintaining the semantic meaning of the existing dimensions, which is particularly useful for broadcasting and reshaping operations in data analysis workflows.

AllenDowney · 2025-05-28T19:47:26Z

@ricardoV94 I've added an implementation of Squeeze. Please take a look.

ricardoV94 · 2025-05-28T19:50:15Z

pytensor/xtensor/shape.py

+
+        # Create new shape with the new dimension
+        new_shape = list(x.type.shape)
+        new_shape.append(1)


I think this is wrong, new dims by default are added on the left, but actually xarray allows specifying a non-default axis: https://docs.xarray.dev/en/stable/generated/xarray.DataArray.expand_dims.html

ricardoV94 · 2025-05-28T19:51:22Z

pytensor/xtensor/shape.py

+            if self.dim not in x.type.dims:
+                raise ValueError(f"Dimension {self.dim} not found")
+            dim_idx = x.type.dims.index(self.dim)
+            if x.type.shape[dim_idx] != 1:


x.type.shape may also be not known None in which case it can be 1 or something else, it's fine to let it through in that case (it will fail at runtime if nto 1). Only if x.type.shape is not None and not 1 do we know it must be invalid

ricardoV94 · 2025-05-28T19:52:25Z

tests/xtensor/test_shape.py

+    xr_assert_allclose(fn(x_test), x_test.transpose())
+
+
+def test_expand_dims():


If you had kept the template of comparing against xarray like the other tests, you would have noticed expand_dims is not doing the same as xarray

Yes, comparing against xarray makes perfect sense. I will do that in the next iteration.

ricardoV94 · 2025-05-28T19:53:04Z

tests/xtensor/test_shape.py

+    xr_assert_allclose(res_multi, expected_res_multi)
+
+
+def test_lower_squeeze():


This one is nonsensical

Cursor and I put that in while debugging -- and it did help. But if it doesn't make sense to keep it as a test, I'll remove it.

ricardoV94 · 2025-05-29T14:21:05Z

@AllenDowney I picked and cleanup up the transpose related changes into #1430.

Regarding squeeze and expand_dims, which I suggest you tackle on a separate PR. Two notes:

Use pytensor.tensor.squeeze / expand_dims instead of reshape. The reason for this is that reshape is a very general Op that is harder to reason symbolically about. We have many rewrites for squeeze/expand_dims, that don't apply to reshape because it's not always easy to know when the two are equivalent
Note that xarray expand_dims allows specifying the size of the new dimension, it need not be length 1. You can do DataArray(np.zeros((2, 3)), dims=("a", "b")).expand_dims(c=5) . This makes sense because xarray does not broadcast existing dims, so they should have the "right" sizes when they are created. This will require a mix of expand_dims and broadcast_to when lowered. It also means expand_dims will be a bit more complex to allow the size of the expand_dims to be symbolic, and will be a bit more like what we are exploring in Add unstack for xtensors #1412. Feel free to split Squeeze from ExpandDims, since the former is simpler.

ricardoV94 and others added 6 commits May 27, 2025 11:06

WIP Implement index operations for XTensorVariables

25397c8

Add diff method to XTensorVariable

e32d865

Add transpose operation for labeled tensors with ellipsis support

5988cec

Refactor: Extract ellipsis expansion logic into helper function

5936ab2

Fix lint errors: remove trailing whitespace from docstrings

6fc7b89

Format files with ruff

0778cf7

Remove commented out line

c7ce0c9

ricardoV94 requested changes May 27, 2025

View reviewed changes

ricardoV94 reviewed May 27, 2025

View reviewed changes

AllenDowney added 7 commits May 28, 2025 10:27

Add missing_dims parameter to transpose for XTensorVariable and core,…

bc2cbc0

… matching xarray behavior; update tests and rewrites accordingly

Add missing_dims parameter to transpose for XTensorVariable and core,…

7bfa2b2

… matching xarray behavior; update tests and rewrites accordingly

Fix linting issues: remove unused Union import and use dict.fromkeys()

d4f5512

Improve expand_ellipsis with validate parameter and update tests

1ed01c4

Apply ruff-format to shape.py, type.py, and test_shape.py for consist…

4f010e0

…ent formatting

Simplify make_node in Transpose class by combining ignore/warn cases

f0ea583

Format expand_ellipsis call for better readability

0125bd2

ricardoV94 force-pushed the labeled_tensors branch from e32d865 to d8fe0d1 Compare May 28, 2025 15:24

ricardoV94 reviewed May 28, 2025

View reviewed changes

ricardoV94 added 2 commits May 28, 2025 19:50

WIP Implement index operations for XTensorVariables

30e1a42

Add diff method to XTensorVariable

29b954a

ricardoV94 force-pushed the labeled_tensors branch from d8fe0d1 to 29b954a Compare May 28, 2025 17:50

AllenDowney added 3 commits May 28, 2025 14:40

Format and simplify expand_ellipsis; auto-fix with pre-commit; update…

a76b15e

… tests

Improve expand_dims: add tests, fix reshape usage, and ensure code st…

af14c90

…yle compliance

Merge WIP changes from origin/labeled_tensors

6208092

ricardoV94 force-pushed the labeled_tensors branch from 29b954a to 5a7b23c Compare May 28, 2025 19:46

Implement squeeze

15f4c48

ricardoV94 reviewed May 28, 2025

View reviewed changes

ricardoV94 mentioned this pull request May 29, 2025

Implement Transpose for XTensorVariables #1430

Merged

ricardoV94 closed this Jun 4, 2025

	# def swap_dims(self, args, *kwargs):
	# ...
	#
	# def expand_dims(self, args, *kwargs):
	# ...
	#
	# def squeeze(self):
	# ...

+                  # Determine the permutation of axes
+                  out_dims = node.op.dims
+                  in_dims = x.type.dims
+                  expanded_dims = expand_ellipsis(out_dims, in_dims)

+              class Transpose(XOp):
+                  __props__ = ("dims",)
+                  def __init__(self, dims: tuple[str, ...]):

	def __init__(self, dims: tuple[str, ...]):
	def __init__(self, dims: tuple[str \| Ellipsis, ...]):

	def isel(
	self,
	indexers: dict[str, Any] \| None = None,
	drop: bool = False, # Unused by PyTensor
	missing_dims: Literal["raise", "warn", "ignore"] = "raise",
	**indexers_kwargs,
	):

		@@ -73,6 +73,67 @@ def stack(x, dim: dict[str, Sequence[str]] \| None = None, **dims: Sequence[str])
		return y


		def expand_ellipsis(

		xr_assert_allclose(fn(x_test), x_test.transpose())


		def test_expand_dims():

		xr_assert_allclose(res_multi, expected_res_multi)


		def test_lower_squeeze():



		def transpose(x, *dims):
		return Transpose(dims)(x)

Add transpose() for labeled tensors #1427

Add transpose() for labeled tensors #1427

Uh oh!

Conversation

AllenDowney commented May 27, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Add transpose operation for labeled tensors

Changes

Features

Testing

Related

Uh oh!

AllenDowney commented May 27, 2025

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 May 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AllenDowney commented May 28, 2025

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

AllenDowney commented May 28, 2025

Uh oh!

AllenDowney commented May 28, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented May 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

AllenDowney commented May 27, 2025 •

edited by github-actions bot

Loading

ricardoV94 May 27, 2025 •

edited

Loading

ricardoV94 commented May 29, 2025 •

edited

Loading