Adding expand_dims for xtensor #1449

AllenDowney · 2025-06-06T18:17:37Z

Add expand_dims operation for labeled tensors

This PR adds support for the expand_dims operation in PyTensor's labeled tensor system, allowing users to add new dimensions to labeled tensors with explicit dimension names.

Key Features

New ExpandDims operation that adds a new dimension to an XTensorVariable
Support for both static and symbolic dimension sizes
Automatic broadcasting when size > 1
Integration with existing tensor operations
Full compatibility with xarray's expand_dims behavior

Implementation Details

The implementation includes:

New ExpandDims class in pytensor/xtensor/shape.py that handles:
- Adding new dimensions with specified names
- Support for both static and symbolic sizes
- Shape inference and validation
Rewriting rule in pytensor/xtensor/rewriting/shape.py that:
- Converts labeled tensor operations to standard tensor operations
- Handles broadcasting when needed
- Validates symbolic sizes
Comprehensive test suite in tests/xtensor/test_shape.py covering:
- Basic dimension expansion
- Static and symbolic sizes
- Error cases and edge cases
- Compatibility with xarray operations
- Integration with other labeled tensor operations

Usage Example

import pytensor.tensor as pt
from pytensor.xtensor import xtensor

# Create a labeled tensor
x = xtensor("x", dims=("city",), shape=(3,))

# Add a new dimension
y = expand_dims(x, "country")  # Adds a new dimension of size 1
z = expand_dims(x, "country", size=4)  # Adds a new dimension of size 4

Testing

The implementation includes extensive tests that verify:

Correct behavior with various input shapes
Proper handling of symbolic sizes
Error cases (invalid dimensions, sizes, etc.)
Compatibility with xarray's expand_dims
Integration with other labeled tensor operations

📚 Documentation preview 📚: https://pytensor--1449.org.readthedocs.build/en/1449/

AllenDowney · 2025-06-06T20:14:55Z

Now that we have this PR based on the right commit, @ricardoV94 it is ready for a first look.

One question: my first draft of this was based on a later commit -- this draft goes back to an earlier commit, and it looks like @register_xcanonicalize doesn't exist yet, so I've replaced it with @register_lower_xtensor, which seems to be its predecessor. Is that the right thing to do for now?

ricardoV94 · 2025-06-06T21:26:26Z

That's the new name, it better represents the kind of rewrites it holds

ricardoV94 · 2025-06-06T21:28:34Z

pytensor/xtensor/shape.py

+
+    def __init__(self, dim, size=1):
+        self.dims = dim
+        self.size = size


This does not allow symbolic sizes, check UnStack for reference

AllenDowney · 2025-06-07T20:12:19Z

@ricardoV94 I think this is a step toward handling symbolic sizes, but there are a couple of place where I'm not sure what the right behavior is. See the comments in test_shape.py, test_expand_dims_implicit.

Do those tests make sense? Are there more cases that should be covered?

ricardoV94 · 2025-06-08T20:40:30Z

The simplest test for symbolic expand_dims is:

size_new_dim = xtensor("size_new_dim", shape=(), dtype=int)
x = xtensor("x", shape=(3,))
y =  x.expand_dims(new_dim=size_new_dim)
xr_function = function([x, size_new_dim], y)

x_test = xr_arange_like(x)
size_new_dim_test = DataArray(np.array(5, dtype=int))
result = xr_function(x_test, size_new_dim_test)
expected_result = x_test.expand_dims(new_dim=size_new_dim_test)
xr_assert_allclose(result, expected_result)

Yout can parametrize the test to try default and explicit non-default axis as well.

Sidenote, what is an implicit expand_dims? I don't think that's a thing.

tests/xtensor/test_shape.py

pytensor/xtensor/shape.py

tests/xtensor/test_shape.py

ricardoV94 · 2025-06-09T06:46:06Z

tests/xtensor/test_shape.py

+
+    # Duplicate dimension creation
+    y = expand_dims(x, "new")
+    with pytest.raises(ValueError):


match the expected error messages to be sure you are triggering the branch you care about. Sometimes you are testing an earlier error and can't tell because of only checking for ValueError/TypeError

AllenDowney · 2025-06-09T15:23:22Z

@ricardoV94 I've addressed most of your comments on the previous round, and made a first pass at adding support for multiple dimensions. Please take a look at the expand_dims wrapper function, which canonicalizes the inputs and loops through them to make a series of Ops.

Assuming that adding multiple dimensions is rare, what do with think of the loop option, as opposed to making a single Op that adds multiple dimensions?

ricardoV94 · 2025-06-09T20:12:19Z

pytensor/xtensor/shape.py

+    def make_node(self, x, size):
+        x = as_xtensor(x)
+
+        if not isinstance(self.dim, str):


This can be done already in init

ricardoV94 · 2025-06-09T20:15:24Z

pytensor/xtensor/shape.py

+        if isinstance(size, int | np.integer):
+            if size <= 0:
+                raise ValueError(f"size must be positive, got: {size}")
+        elif not (
+            hasattr(size, "ndim")
+            and getattr(size, "ndim", None) == 0  # symbolic scalar
+        ):
+            raise TypeError(
+                f"size must be an int or scalar variable, got: {type(size)}"
+            )
+
+        # Convert size to tensor
+        size = as_tensor(size, ndim=0)


Better to convert to as_xtensor, since both tensor and xtensor scalars will be accepted then.

Instead of having the check for numpy specifically, check if after conversion you have isinstance(size, Constant) and (size.data <= 0) and then raise. That will cover more cases. Or actually use the inferred static size variable below to do that check.

Are you sure xarray doesn't support size of zero? It's valid otherwise in numpy

ricardoV94 · 2025-06-09T20:16:00Z

pytensor/xtensor/rewriting/shape.py

+    x, size = node.inputs
+    out = node.outputs[0]
+    # Lower to tensor.expand_dims(x, axis=0)
+    from pytensor.tensor import expand_dims as tensor_expand_dims


import globally

ricardoV94 · 2025-06-09T20:17:11Z

pytensor/xtensor/rewriting/shape.py

+    from pytensor.tensor import broadcast_to
+
+    # Ensure size is positive
+    expanded = Assert(msg="size must be positive")(expanded, gt(size, 0))


No need for assert which pollutes the graph, broadcast_to will fail at runtime if size is negative. Size zero is allowed but if it's also allowed in xarray that's fine

ricardoV94 · 2025-06-09T20:17:59Z

pytensor/xtensor/rewriting/shape.py

+    expanded = Assert(msg="size must be positive")(expanded, gt(size, 0))
+    # If size is not 1, broadcast
+    try:
+        static_size = get_scalar_constant_value(size)


You already inferred static size in make_node so you can use that info from out.type.shape

ricardoV94 · 2025-06-09T20:21:16Z

tests/xtensor/test_shape.py

+
+    # Test behavior with symbolic size > 1
+    # NOTE: This test documents our current behavior where expand_dims broadcasts to the requested size.
+    # This differs from xarray's behavior where expand_dims always adds a size-1 dimension.


This is not true?

ricardoV94 · 2025-06-09T20:23:00Z

Assuming that adding multiple dimensions is rare, what do with think of the loop option, as opposed to making a single Op that adds multiple dimensions?

That's fine. We used that for other Ops and we can revisit later of we want it to be fused

Adding expand_dims for xtensor

19361d6

ricardoV94 reviewed Jun 6, 2025

View reviewed changes

Handling symbolic size

95a71a6

ricardoV94 reviewed Jun 8, 2025

View reviewed changes

tests/xtensor/test_shape.py Show resolved Hide resolved

ricardoV94 reviewed Jun 9, 2025

View reviewed changes

Adding support for expanding multiple dimensions

802a536

ricardoV94 reviewed Jun 9, 2025

View reviewed changes

Adding expand_dims for xtensor #1449

Are you sure you want to change the base?

Adding expand_dims for xtensor #1449

Uh oh!

Conversation

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

Uh oh!

Add expand_dims operation for labeled tensors

Key Features

Implementation Details

Usage Example

Testing

Uh oh!

AllenDowney commented Jun 6, 2025

Uh oh!

ricardoV94 commented Jun 6, 2025

Uh oh!

ricardoV94 Jun 6, 2025

Choose a reason for hiding this comment

Uh oh!

AllenDowney commented Jun 7, 2025

Uh oh!

ricardoV94 commented Jun 8, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ricardoV94 Jun 9, 2025

Choose a reason for hiding this comment

Uh oh!

AllenDowney commented Jun 9, 2025

Uh oh!

ricardoV94 Jun 9, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jun 9, 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!

ricardoV94 Jun 9, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jun 9, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jun 9, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jun 9, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented Jun 9, 2025

Uh oh!

Uh oh!

AllenDowney commented Jun 6, 2025 •

edited by github-actions bot

Loading

ricardoV94 Jun 9, 2025 •

edited

Loading