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

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

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

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.

[ricardoV94]

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.

[AllenDowney]

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

[AllenDowney]

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

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

[AllenDowney]

Right, but True is the default. What do you think of setting the default value to None and warning if it's set to anything else?

[ricardoV94]

I think that's good

[ricardoV94]

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

[AllenDowney]

@ricardoV94 This is ready for another look