Adding expand_dims for xtensor

[AllenDowney]

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:

1. 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

2. Rewriting rule in pytensor/xtensor/rewriting/shape.py that:

   • Converts labeled tensor operations to standard tensor operations
   • Handles broadcasting when needed
   • Validates symbolic sizes

3. 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]

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]

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

[ricardoV94]

This does not allow symbolic sizes, check UnStack for reference

[AllenDowney]

@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]

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.

[ricardoV94]

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]

@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]

This is not true?

[AllenDowney]

I'm not sure about the general claim in the note, but at least in this case, it seems like we're getting the behavior we want from xtensor, but running the same operation with xarray does something different, causing the test to fail. Here's Cursor's summary

The test failure confirms that our current implementation of expand_dims broadcasts to the requested size (4 in this case), while xarray's behavior is to always add a size-1 dimension. This is evident from the test output, where the left side (our implementation) has a shape of (batch: 4, a: 2, b: 3), and the right side (xarray's behavior) has a shape of (batch: 1, a: 2, b: 3).

I'm inclined to keep this test to note the discrepancy.

[ricardoV94]

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]

@ricardoV94 This is ready for another look.

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