Cleanup

Author: AllenDowney

pytensor/xtensor/rewriting/shape.py

@@ -123,7 +123,7 @@ def lower_transpose(fgraph, node):
 
 @register_lower_xtensor
 @node_rewriter([Squeeze])
-def local_squeeze_reshape(fgraph, node):
+def lower_squeeze(fgraph, node):
     """Rewrite Squeeze to tensor.squeeze."""
     [x] = node.inputs
     x_tensor = tensor_from_xtensor(x)
@@ -138,7 +138,7 @@ def local_squeeze_reshape(fgraph, node):
 
 @register_lower_xtensor
 @node_rewriter([ExpandDims])
-def local_expand_dims_reshape(fgraph, node):
+def lower_expand_dims(fgraph, node):
     """Rewrite ExpandDims using tensor operations."""
     x, size = node.inputs
     out = node.outputs[0]
@@ -155,10 +155,8 @@ def local_expand_dims_reshape(fgraph, node):
         # Simple case: just expand with size 1
         result_tensor = expand_dims(x_tensor, new_axis)
     else:
-        # First expand with size 1
-        expanded = expand_dims(x_tensor, new_axis)
-        # Then broadcast to the requested size
-        result_tensor = broadcast_to(expanded, (size_tensor, *x_tensor.shape))
+        # Otherwise broadcast to the requested size
+        result_tensor = broadcast_to(x_tensor, (size_tensor, *x_tensor.shape))
 
     # Preserve static shape information
     result_tensor = specify_shape(result_tensor, out.type.shape)

pytensor/xtensor/shape.py

@@ -442,32 +442,16 @@ def make_node(self, x, size):
 
 
 def expand_dims(x, dim=None, create_index_for_new_dim=True, **dim_kwargs):
-    """Add one or more new dimensions to an XTensorVariable.
-
-    Parameters
-    ----------
-    x : XTensorVariable
-        Input tensor.
-    dim : str | Sequence[str] | dict[str, int | Sequence] | None
-        If str or sequence of str, new dimensions with size 1.
-        If dict, keys are dimension names and values are either:
-            - 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.
-    **dim_kwargs : int | Sequence
-        Alternative to `dim` dict. Only used if `dim` is None.
-
-    Returns
-    -------
-    XTensorVariable
-        A tensor with additional dimensions inserted at the front.
-    """
+    """Add one or more new dimensions to an XTensorVariable."""
     x = as_xtensor(x)
 
     # Extract size from dim_kwargs if present
     size = dim_kwargs.pop("size", 1) if dim_kwargs else 1
 
+    # 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")
+
     if dim is None:
         dim = dim_kwargs
     elif dim_kwargs:

pytensor/xtensor/type.py

@@ -483,24 +483,32 @@ def squeeze(
 
     def expand_dims(
         self,
-        dim: str | None = None,
-        size: int | Variable = 1,
+        dim: str | Sequence[str] | dict[str, int | Sequence] | None = None,
+        create_index_for_new_dim: bool = True,
+        **dim_kwargs,
     ):
-        """Add a new dimension to the tensor.
+        """Add one or more new dimensions to the tensor.
 
         Parameters
         ----------
-        dim : str or None
-            Name of new dimension. If None, returns self unchanged.
-        size : int or symbolic, optional
-            Size of the new dimension (default 1)
+        dim : str | Sequence[str] | dict[str, int | Sequence] | None
+            If str or sequence of str, new dimensions with size 1.
+            If dict, keys are dimension names and values are either:
+                - 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.
+        **dim_kwargs : int | Sequence
+            Alternative to `dim` dict. Only used if `dim` is None.
 
         Returns
         -------
         XTensorVariable
-            Tensor with the new dimension inserted
+            A tensor with additional dimensions inserted at the front.
         """
-        return px.shape.expand_dims(self, dim, size=size)
+        return px.shape.expand_dims(
+            self, dim, create_index_for_new_dim=create_index_for_new_dim, **dim_kwargs
+        )
 
     # ndarray methods
     # https://docs.xarray.dev/en/latest/api.html#id7

tests/xtensor/test_shape.py

@@ -8,14 +8,12 @@
 from itertools import chain, combinations
 
 import numpy as np
-import xarray as xr
 from xarray import DataArray
 from xarray import concat as xr_concat
 
 from pytensor.tensor import scalar
 from pytensor.xtensor.shape import (
     concat,
-    expand_dims,
     squeeze,
     stack,
     transpose,
@@ -373,211 +371,97 @@ def test_squeeze_errors():
         fn2(x2_test)
 
 
-def test_expand_dims_explicit():
-    """Test expand_dims with explicitly named dimensions and sizes."""
-
-    # 1D case
-    x = xtensor("x", dims=("city",), shape=(3,))
-    y = expand_dims(x, "country")
-    fn = xr_function([x], y)
-    x_xr = xr_arange_like(x)
-    xr_assert_allclose(fn(x_xr), x_xr.expand_dims("country"))
-
-    # 2D case
+def test_expand_dims():
+    """Test expand_dims."""
     x = xtensor("x", dims=("city", "year"), shape=(2, 2))
-    y = expand_dims(x, "country")
-    fn = xr_function([x], y)
-    xr_assert_allclose(fn(xr_arange_like(x)), xr_arange_like(x).expand_dims("country"))
-
-    # 3D case
-    x = xtensor("x", dims=("city", "year", "month"), shape=(2, 2, 2))
-    y = expand_dims(x, "country")
-    fn = xr_function([x], y)
-    xr_assert_allclose(fn(xr_arange_like(x)), xr_arange_like(x).expand_dims("country"))
-
-    # Prepending various dims
-    x = xtensor("x", dims=("a", "b"), shape=(2, 3))
-    for new_dim in ("x", "y", "z"):
-        y = expand_dims(x, new_dim)
-        assert y.type.dims == (new_dim, "a", "b")
-        assert y.type.shape == (1, 2, 3)
-
-    # Explicit size=1 behaves like default
-    y1 = expand_dims(x, "batch", size=1)
-    y2 = expand_dims(x, "batch")
-    fn1 = xr_function([x], y1)
-    fn2 = xr_function([x], y2)
     x_test = xr_arange_like(x)
-    xr_assert_allclose(fn1(x_test), fn2(x_test))
 
-    # Scalar expansion
-    x = xtensor("x", dims=(), shape=())
-    y = expand_dims(x, "batch")
-    assert y.type.dims == ("batch",)
-    assert y.type.shape == (1,)
+    # Implicit size=1
+    y = x.expand_dims("country")
     fn = xr_function([x], y)
-    xr_assert_allclose(fn(xr_arange_like(x)), xr_arange_like(x).expand_dims("batch"))
+    xr_assert_allclose(fn(x_test), x_test.expand_dims("country"))
 
-    # Static size > 1: broadcast
-    x = xtensor("x", dims=("a", "b"), shape=(2, 3))
-    y = expand_dims(x, "batch", size=4)
-    fn = xr_function([x], y)
-    expected = xr.DataArray(
-        np.broadcast_to(xr_arange_like(x).data, (4, 2, 3)),
-        dims=("batch", "a", "b"),
-        coords={"a": xr_arange_like(x).coords["a"], "b": xr_arange_like(x).coords["b"]},
-    )
-    xr_assert_allclose(fn(xr_arange_like(x)), expected)
+    # Explicit size=1
+    y = x.expand_dims("country", size=1)
+    xr_assert_allclose(fn(x_test), x_test.expand_dims("country"))
 
-    # Insert new dim between existing dims
-    x = xtensor("x", dims=("a", "b"), shape=(2, 3))
-    y = expand_dims(x, "new")
-    # Insert new dim between a and b: ("a", "new", "b")
-    y = transpose(y, "a", "new", "b")
+    # Explicit size > 1
+    y = x.expand_dims("country", size=4)
     fn = xr_function([x], y)
-    x_test = xr_arange_like(x)
-    expected = x_test.expand_dims("new").transpose("a", "new", "b")
-    xr_assert_allclose(fn(x_test), expected)
+    xr_assert_allclose(fn(x_test), x_test.expand_dims({"country": 4}))
 
-    # Expand with multiple dims
-    x = xtensor("x", dims=(), shape=())
-    y = expand_dims(expand_dims(x, "a"), "b")
+    # Test with multiple dimensions
+    y = x.expand_dims(["country", "state"])
     fn = xr_function([x], y)
-    expected = xr_arange_like(x).expand_dims("a").expand_dims("b")
-    xr_assert_allclose(fn(xr_arange_like(x)), expected)
+    xr_assert_allclose(fn(x_test), x_test.expand_dims(["country", "state"]))
 
+    # Test with a dict of sizes
+    y = x.expand_dims({"country": 2, "state": 3})
+    fn = xr_function([x], y)
+    xr_assert_allclose(fn(x_test), x_test.expand_dims({"country": 2, "state": 3}))
 
-def test_expand_dims_symbolic_size():
-    """Test expand_dims with symbolic sizes."""
-
-    # Symbolic size=1: same as default
-    size_sym_1 = scalar("size_sym_1", dtype="int64")
-    x = xtensor("x", dims=("a", "b"), shape=(2, 3))
-    y = expand_dims(x, "batch", size=size_sym_1)
-    fn = xr_function([x, size_sym_1], y, on_unused_input="ignore")
-    x_test = xr_arange_like(x)
-    xr_assert_allclose(fn(x_test, 1), x_test.expand_dims("batch"))
-
-    # Test using symbolic size from an existing dimension of the same tensor
-    # This verifies that expand_dims can use the size of one dimension to create another
-    x = xtensor(dims=("a", "b", "c"))
-    y = expand_dims(x, "d", size=x.sizes["b"])
+    # Test with kwargs (equivalent to dict)
+    y = x.expand_dims(country=2, state=3)
     fn = xr_function([x], y)
-    x_test = xr_arange_like(xtensor(dims=x.dims, shape=(2, 3, 5)))
-    res = fn(x_test)
-    expected = x_test.expand_dims({"d": 3})  # 3 is the size of dimension "b"
-    xr_assert_allclose(res, expected)
+    xr_assert_allclose(fn(x_test), x_test.expand_dims(country=2, state=3))
 
-    # Test broadcasting with symbolic size from a different tensor
-    x = xtensor("x", dims=("a", "b"), shape=(2, 3))
-    other = xtensor("other", dims=("c",), shape=(4,))
-    y = expand_dims(x, "batch", size=other.sizes["c"])
-    fn = xr_function([x, other], y)
-    x_test = xr_arange_like(x)
-    other_test = xr_arange_like(other)
-    res = fn(x_test, other_test)
-    expected = x_test.expand_dims(
-        {"batch": 4}
-    )  # 4 is the size of dimension "c" in other
-    xr_assert_allclose(res, expected)
+    # Symbolic size=1
+    size_sym_1 = scalar("size_sym_1", dtype="int64")
+    y = x.expand_dims("country", size=size_sym_1)
+    fn = xr_function([x, size_sym_1], y)
+    xr_assert_allclose(fn(x_test, 1), x_test.expand_dims("country"))
 
     # Test behavior with symbolic size > 1
     # NOTE: This test documents our current behavior where expand_dims broadcasts to the requested size.
     # This differs from xarray's behavior where expand_dims always adds a size-1 dimension.
     size_sym_4 = scalar("size_sym_4", dtype="int64")
-    x = xtensor("x", dims=("a", "b"), shape=(2, 3))
-    y = expand_dims(x, "batch", size=size_sym_4)
-    fn = xr_function([x, size_sym_4], y, on_unused_input="ignore")
-    x_test = xr_arange_like(x)
+    y = x.expand_dims("country", size=size_sym_4)
+    fn = xr_function([x, size_sym_4], y)
     res = fn(x_test, 4)
     # Our current behavior: broadcasts to size 4
-    expected = x_test.expand_dims({"batch": 4})
+    expected = x_test.expand_dims({"country": 4})
     xr_assert_allclose(res, expected)
     # xarray's behavior would be:
-    # expected = x_test.expand_dims("batch")  # always size 1
+    # expected = x_test.expand_dims("country")  # always size 1
     # xr_assert_allclose(res, expected)
 
-    # Test using symbolic size from a reduction operation
-    x = xtensor("x", dims=("a", "b"), shape=(2, 3))
-    reduced = x.sum("a")  # shape: (b: 3)
-    y = expand_dims(x, "batch", size=reduced.sizes["b"])
-    fn = xr_function([x], y)
-    x_test = xr_arange_like(x)
-    res = fn(x_test)
-    expected = x_test.expand_dims({"batch": 3})  # 3 is the size of dimension "b"
-    xr_assert_allclose(res, expected)
-
-    # Test chaining expand_dims with symbolic sizes
-    x = xtensor("x", dims=("a",), shape=(2,))
-    y = expand_dims(x, "b", size=x.sizes["a"])  # shape: (a: 2, b: 2)
-    z = expand_dims(y, "c", size=y.sizes["b"])  # shape: (a: 2, b: 2, c: 2)
-    fn = xr_function([x], z)
-    x_test = xr_arange_like(x)
-    res = fn(x_test)
-    expected = x_test.expand_dims({"b": 2}).expand_dims({"c": 2})
-    xr_assert_allclose(res, expected)
+    # Test with symbolic sizes in dict
+    size_sym_1 = scalar("size_sym_1", dtype="int64")
+    size_sym_2 = scalar("size_sym_2", dtype="int64")
+    y = x.expand_dims({"country": size_sym_1, "state": size_sym_2})
+    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 bidirectional broadcasting with symbolic sizes
-    x = xtensor("x", dims=("a",), shape=(2,))
-    y = xtensor("y", dims=("b",), shape=(3,))
-    # Expand x with size from y, then add y
-    expanded = expand_dims(x, "b", size=y.sizes["b"])
-    z = expanded + y  # Should broadcast x to match y's size
-    fn = xr_function([x, y], z)
-    x_test = xr_arange_like(x)
-    y_test = xr_arange_like(y)
-    res = fn(x_test, y_test)
-    expected = x_test.expand_dims({"b": 3}) + y_test
-    xr_assert_allclose(res, expected)
+    # Test with symbolic sizes in kwargs
+    y = x.expand_dims(country=size_sym_1, state=size_sym_2)
+    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}))
 
 
 def test_expand_dims_errors():
     """Test error handling in expand_dims."""
 
     # Expanding existing dim
     x = xtensor("x", dims=("city",), shape=(3,))
-    y = expand_dims(x, "country")
+    y = x.expand_dims("country")
     with pytest.raises(ValueError, match="already exists"):
-        expand_dims(y, "city")
+        y.expand_dims("city")
 
     # Invalid dim type
     with pytest.raises(TypeError, match="Invalid type for `dim`"):
-        expand_dims(x, 123)
+        x.expand_dims(123)
 
     # Invalid size type
     with pytest.raises(TypeError, match="size must be an int or scalar variable"):
-        expand_dims(x, "new", size=[1])
+        x.expand_dims("new", size=[1])
 
     # Duplicate dimension creation
-    y = expand_dims(x, "new")
+    y = x.expand_dims("new")
     with pytest.raises(ValueError, match="already exists"):
-        expand_dims(y, "new")
-
+        y.expand_dims("new")
 
-def test_expand_dims_multiple():
-    """Test expanding multiple dimensions at once using a list of strings."""
-    x = xtensor("x", dims=("city",), shape=(3,))
-    y = expand_dims(x, ["country", "state"])
-    fn = xr_function([x], y)
-    x_xr = xr_arange_like(x)
-    xr_assert_allclose(fn(x_xr), x_xr.expand_dims(["country", "state"]))
-
-    # Test with a dict of sizes
-    y = expand_dims(x, {"country": 2, "state": 3})
-    fn = xr_function([x], y)
-    x_xr = xr_arange_like(x)
-    xr_assert_allclose(fn(x_xr), x_xr.expand_dims({"country": 2, "state": 3}))
-
-    # Test with a mix of strings and dicts
-    y = expand_dims(x, ["country", "state"], size=3)
-    fn = xr_function([x], y)
-    x_xr = xr_arange_like(x)
-    xr_assert_allclose(fn(x_xr), x_xr.expand_dims(["country", "state"]))
-
-    # Test with symbolic sizes in dict
-    size_sym_1 = scalar("size_sym_1", dtype="int64")
-    size_sym_2 = scalar("size_sym_2", dtype="int64")
-    y = expand_dims(x, {"country": size_sym_1, "state": size_sym_2})
-    fn = xr_function([x, size_sym_1, size_sym_2], y, on_unused_input="ignore")
-    x_xr = xr_arange_like(x)
-    xr_assert_allclose(fn(x_xr, 2, 3), x_xr.expand_dims({"country": 2, "state": 3}))
+    # Test for error when both positional and size are given
+    with pytest.raises(
+        ValueError, match="cannot specify both keyword and positional arguments"
+    ):
+        x.expand_dims(["country", "state"], size=3)