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?