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

pytensor/xtensor/shape.py

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

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?

pytensor/xtensor/shape.py

pytensor/xtensor/rewriting/shape.py

tests/xtensor/test_shape.py

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

AllenDowney · 2025-06-10T16:22:38Z

@ricardoV94 This is ready for another look.

The rewrite was a shambles, but I think I have a clearer idea now.

tests/xtensor/test_shape.py

pytensor/xtensor/rewriting/shape.py

ricardoV94 · 2025-06-10T22:27:19Z

I left some comments above.

Rewrite looks good. As we discussed we should redo the tests to use expand_dims as a method (like xarray users would).

Also, I suspected xarray allows specifying the size like x.expand_dims(dim_a=1, dim_b=2) which is equivalent to x.expand_dims({"dim_a":1, "dim_b":2}). At least that was a pattern I noticed in other xarray methods. I saw you had a test for multiple dims with dict, but I didn't see one with kwargs.

pytensor/xtensor/rewriting/shape.py

pytensor/xtensor/shape.py

tests/xtensor/test_shape.py

AllenDowney · 2025-06-11T19:08:44Z

@ricardoV94 I cleaned up the code as suggested and took a first cut at handling the axis parameter. Please take a look.

ricardoV94

This looks great. Just the question about the size kwarg and small notes.

pytensor/xtensor/shape.py

pytensor/xtensor/type.py

tests/xtensor/test_shape.py

pytensor/xtensor/shape.py

AllenDowney · 2025-06-12T14:46:25Z

@ricardoV94 This is ready for another look

ricardoV94

I think we're nearly there. This proved to be more complex than I antecipated.

Small error checks changes and one question. Also I don't think there's any test for the passing sequences as the length of the dimensions?

ricardoV94 · 2025-06-12T17:16:21Z

pytensor/xtensor/shape.py

+    if not create_index_for_new_dim:
+        warnings.warn(
+            "create_index_for_new_dim=False has no effect in pytensor.xtensor",
+            UserWarning,
+            stacklevel=2,
+        )


Neither option has an effect to be fair, just don't mention?

ricardoV94 · 2025-06-12T17:16:39Z

pytensor/xtensor/shape.py

+        )
+
+    # Extract size from dim_kwargs if present
+    size = dim_kwargs.pop("size", 1) if dim_kwargs else 1


Still here, did you not push yet?

ricardoV94 · 2025-06-12T17:17:12Z

pytensor/xtensor/shape.py

+    # xarray compatibility: error if a sequence (list/tuple) of dims and size are given
+    if (isinstance(dim, list | tuple)) and ("size" in locals() and size != 1):
+        raise ValueError("cannot specify both keyword and positional arguments")


also not a thing?

ricardoV94 · 2025-06-12T17:17:23Z

pytensor/xtensor/shape.py

+
+    # Normalize to a dimension-size mapping
+    if isinstance(dim, str):
+        dims_dict = {dim: size}


Suggested change

dims_dict = {dim: size}

dims_dict = {dim: 1}

ricardoV94 · 2025-06-12T17:23:16Z

pytensor/xtensor/shape.py

+    elif isinstance(dim, dict):
+        dims_dict = {}
+        for name, val in dim.items():
+            if isinstance(val, Sequence | np.ndarray) and not isinstance(val, str):


How does xarray treat expand_dims(new_dim=np.array(5))? Does it treat as coordinates or the size? I suppose the latter?

In that case the check here should be (isinstance(val, np.ndarray), and val.ndim > 0). We could also consider symbolic variables with (isinstance(val, np.ndarray) or (isinstance(val, Variable) and isinstance(val.type, HasShape)) and val.ndim > 0)

ricardoV94 · 2025-06-12T17:24:18Z

pytensor/xtensor/shape.py

+            elif isinstance(val, int):
+                dims_dict[name] = val
+            else:
+                dims_dict[name] = val  # symbolic/int scalar allowed


This ends up being accepted anyway so merge?

Suggested change

elif isinstance(val, int):

dims_dict[name] = val

else:

dims_dict[name] = val # symbolic/int scalar allowed

else:

dims_dict[name] = val # symbolic/int scalar allowed

ricardoV94 · 2025-06-12T17:25:24Z

pytensor/xtensor/shape.py

+    elif isinstance(dim, dict):
+        dims_dict = {}
+        for name, val in dim.items():
+            if isinstance(val, Sequence | np.ndarray) and not isinstance(val, str):


Suggested change

if isinstance(val, Sequence | np.ndarray) and not isinstance(val, str):

if isinstance(val, str):

raise ValueError(f"The size of a dimension cannot be a string, got {val})

if isinstance(val, Sequence | np.ndarray):

ricardoV94 · 2025-06-12T17:27:53Z

PS it's so annoying there's no type for non-string sequences

Adding expand_dims for xtensor

19361d6

ricardoV94 reviewed Jun 6, 2025

View reviewed changes

pytensor/xtensor/shape.py Outdated Show resolved Hide resolved

Handling symbolic size

95a71a6

ricardoV94 reviewed Jun 8, 2025

View reviewed changes

tests/xtensor/test_shape.py Outdated 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

Fixing ExpandDims rewrite

3b7b973

ricardoV94 requested changes Jun 10, 2025

View reviewed changes

ricardoV94 reviewed Jun 10, 2025

View reviewed changes

pytensor/xtensor/rewriting/shape.py Outdated Show resolved Hide resolved

pytensor/xtensor/shape.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Jun 10, 2025

View reviewed changes

tests/xtensor/test_shape.py Outdated Show resolved Hide resolved

AllenDowney added 2 commits June 11, 2025 14:44

Cleanup

38d87c6

Adding axis parameter

30354cd

AllenDowney added 2 commits June 11, 2025 15:20

Fixing axis parameter

efc0fdb

Clean up make_node

3d15197

ricardoV94 reviewed Jun 11, 2025

View reviewed changes

Cleanup

6f810d3

Cleanup

841e52f

ricardoV94 reviewed Jun 12, 2025

View reviewed changes

-            if isinstance(val, Sequence | np.ndarray) and not isinstance(val, str):
+            if isinstance(val, str):
+                raise ValueError(f"The size of a dimension cannot be a string, got {val})
+            if isinstance(val, Sequence | np.ndarray):

Adding expand_dims for xtensor #1449

Are you sure you want to change the base?

Adding expand_dims for xtensor #1449

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!

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!

Uh oh!

AllenDowney commented Jun 9, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ricardoV94 commented Jun 9, 2025

Uh oh!

AllenDowney commented Jun 10, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

AllenDowney commented Jun 11, 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!

Uh oh!

Uh oh!

Uh oh!

AllenDowney commented Jun 12, 2025

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jun 12, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jun 12, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jun 12, 2025

AllenDowney commented Jun 6, 2025 •

edited by github-actions bot

Loading

ricardoV94 commented Jun 10, 2025 •

edited

Loading