Adding axis parameter

Author: AllenDowney

pytensor/xtensor/shape.py

@@ -441,7 +441,7 @@ def make_node(self, x, size):
         return Apply(self, [x, size], [out])
 
 
-def expand_dims(x, dim=None, create_index_for_new_dim=True, **dim_kwargs):
+def expand_dims(x, dim=None, create_index_for_new_dim=True, axis=None, **dim_kwargs):
     """Add one or more new dimensions to an XTensorVariable."""
     x = as_xtensor(x)
 
@@ -479,8 +479,32 @@ def expand_dims(x, dim=None, create_index_for_new_dim=True, **dim_kwargs):
     for name, size in dims_dict.items():
         canonical_dims.append((name, size))
 
+    # Store original dimensions for later use with axis
+    original_dims = list(x.type.dims)
+
     # Insert each new dim at the front (reverse order preserves user intent)
     for name, size in reversed(canonical_dims):
         x = ExpandDims(dim=name)(x, size)
 
+    # If axis is specified, transpose to put new dimensions in the right place
+    if axis is not None:
+        new_dim_names = [name for name, _ in canonical_dims]
+        # Wrap non-sequence axis in a list
+        if not isinstance(axis, Sequence):
+            axis = [axis]
+
+        # xarray requires len(axis) == len(new_dim_names)
+        if len(axis) != len(new_dim_names):
+            raise ValueError("lengths of dim and axis should be identical.")
+
+        # Insert each new dim at the specified axis position
+        # Start with original dims, then insert each new dim at its axis
+        target_dims = list(original_dims)
+        # axis values are relative to the result after each insertion
+        for insert_dim, insert_axis in sorted(
+            zip(new_dim_names, axis), key=lambda x: x[1]
+        ):
+            target_dims.insert(insert_axis, insert_dim)
+        x = transpose(x, *target_dims)
+
     return x

pytensor/xtensor/type.py

@@ -485,6 +485,7 @@ def expand_dims(
         self,
         dim: str | Sequence[str] | dict[str, int | Sequence] | None = None,
         create_index_for_new_dim: bool = True,
+        axis: int | None = None,
         **dim_kwargs,
     ):
         """Add one or more new dimensions to the tensor.
@@ -497,7 +498,14 @@ def expand_dims(
                 - int: the new size
                 - sequence: coordinates (length determines size)
         create_index_for_new_dim : bool, default: True
-            (Ignored for now) Matches xarray API, reserved for future use.
+            Currently ignored. Reserved for future coordinate support.
+            In xarray, when True (default), creates a coordinate index for the new dimension
+            with values from 0 to size-1. When False, no coordinate index is created.
+        axis : int | None, default: None
+            Not implemented yet. In xarray, specifies where to insert the new dimension(s).
+            By default (None), new dimensions are inserted at the beginning (axis=0).
+            Symbolic axis is not supported yet.
+            Negative values count from the end.
         **dim_kwargs : int | Sequence
             Alternative to `dim` dict. Only used if `dim` is None.
 
@@ -507,7 +515,11 @@ def expand_dims(
             A tensor with additional dimensions inserted at the front.
         """
         return px.shape.expand_dims(
-            self, dim, create_index_for_new_dim=create_index_for_new_dim, **dim_kwargs
+            self,
+            dim,
+            create_index_for_new_dim=create_index_for_new_dim,
+            axis=axis,
+            **dim_kwargs,
         )
 
     # ndarray methods

tests/xtensor/test_shape.py

@@ -437,6 +437,18 @@ def test_expand_dims():
     fn = xr_function([x, size_sym_1, size_sym_2], y)
     xr_assert_allclose(fn(x_test, 2, 3), x_test.expand_dims({"country": 2, "state": 3}))
 
+    # Test with axis parameter
+    y = x.expand_dims("country", axis=1)
+    fn = xr_function([x], y)
+    xr_assert_allclose(fn(x_test), x_test.expand_dims("country", axis=1))
+
+    # Add two new dims at axis=[1, 2]
+    y = x.expand_dims(["country", "state"], axis=[1, 2])
+    fn = xr_function([x], y)
+    xr_assert_allclose(
+        fn(x_test), x_test.expand_dims(["country", "state"], axis=[1, 2])
+    )
+
 
 def test_expand_dims_errors():
     """Test error handling in expand_dims."""