Removing expand_dims

Author: AllenDowney

pytensor/xtensor/rewriting/shape.py

@@ -1,11 +1,6 @@
-import numpy as np
-
 from pytensor.graph import node_rewriter
-from pytensor.raise_op import Assert
 from pytensor.tensor import (
     broadcast_to,
-    expand_dims,
-    gt,
     join,
     moveaxis,
     specify_shape,
@@ -15,7 +10,6 @@
 from pytensor.xtensor.rewriting.basic import register_xcanonicalize
 from pytensor.xtensor.shape import (
     Concat,
-    ExpandDims,
     Squeeze,
     Stack,
     Transpose,
@@ -125,40 +119,6 @@ def lower_transpose(fgraph, node):
     return [new_out]
 
 
-@register_xcanonicalize
-@node_rewriter([ExpandDims])
-def local_expand_dims_reshape(fgraph, node):
-    """Rewrite ExpandDims to tensor.expand_dims and optionally broadcast_to or specify_shape."""
-    if not isinstance(node.op, ExpandDims):
-        return False
-
-    x = node.inputs[0]
-    dim = node.op.dims
-    size = getattr(node.op, "size", 1)
-
-    if dim is None:
-        return [x]
-
-    x_tensor = tensor_from_xtensor(x)
-    x_tensor_expanded = expand_dims(x_tensor, axis=0)
-
-    target_shape = node.outputs[0].type.shape
-
-    if isinstance(size, int | np.integer):
-        if size != 1 and None not in target_shape:
-            x_tensor_expanded = broadcast_to(x_tensor_expanded, target_shape)
-    else:
-        # Symbolic size: enforce shape so broadcast happens downstream correctly
-        # Also validate that size is positive
-        x_tensor_expanded = Assert(msg="size must be positive")(
-            x_tensor_expanded, gt(size, 0)
-        )
-        x_tensor_expanded = specify_shape(x_tensor_expanded, target_shape)
-
-    new_out = xtensor_from_tensor(x_tensor_expanded, dims=node.outputs[0].type.dims)
-    return [new_out]
-
-
 @register_xcanonicalize
 @node_rewriter([Squeeze])
 def local_squeeze_reshape(fgraph, node):

pytensor/xtensor/shape.py

@@ -2,8 +2,6 @@
 from collections.abc import Sequence
 from typing import Literal
 
-import numpy as np
-
 from pytensor import Variable
 from pytensor.graph import Apply
 from pytensor.scalar import discrete_dtypes, upcast
@@ -305,84 +303,6 @@ def concat(xtensors, dim: str):
     return Concat(dim=dim)(*xtensors)
 
 
-class ExpandDims(XOp):
-    """Add a new dimension to an XTensorVariable."""
-
-    __props__ = ("dims", "size")
-
-    def __init__(self, dim, size=1):
-        self.dims = dim
-        self.size = size
-
-    def make_node(self, x):
-        x = as_xtensor(x)
-
-        if self.dims is None:
-            # No-op: return same variable
-            return Apply(self, [x], [x])
-
-        # Insert new dim at front
-        new_dims = (self.dims, *x.type.dims)
-
-        # Determine shape
-        if isinstance(self.size, int | np.integer):
-            new_shape = (self.size, *x.type.shape)
-        else:
-            new_shape = (None, *x.type.shape)  # symbolic size
-
-        out = xtensor(
-            dtype=x.type.dtype,
-            shape=new_shape,
-            dims=new_dims,
-        )
-        return Apply(self, [x], [out])
-
-    def infer_shape(self, fgraph, node, input_shapes):
-        (input_shape,) = input_shapes
-        if self.dims is None:
-            return [input_shape]
-        return [(self.size, *list(input_shape))]
-
-
-def expand_dims(x, dim: str | None, size=1):
-    """Add a new dimension to an XTensorVariable.
-
-    Parameters
-    ----------
-    x : XTensorVariable
-        Input tensor
-    dim : str or None
-        Name of new dimension. If None, returns x unchanged.
-    size : int or symbolic, optional
-        Size of the new dimension (default 1)
-
-    Returns
-    -------
-    XTensorVariable
-        Tensor with the new dimension inserted
-    """
-    x = as_xtensor(x)
-
-    if dim is None:
-        return x  # No-op
-
-    if not isinstance(dim, str):
-        raise TypeError(f"`dim` must be a string or None, got: {type(dim)}")
-
-    if dim in x.type.dims:
-        raise ValueError(f"Dimension {dim} already exists in {x.type.dims}")
-
-    if isinstance(size, int | np.integer):
-        if size <= 0:
-            raise ValueError(f"size must be positive, got: {size}")
-    elif not (
-        hasattr(size, "ndim") and getattr(size, "ndim", None) == 0  # symbolic scalar
-    ):
-        raise TypeError(f"size must be an int or scalar variable, got: {type(size)}")
-
-    return ExpandDims(dim=dim, size=size)(x)
-
-
 class Squeeze(XOp):
     """Remove specified dimensions from an XTensorVariable.
 

tests/xtensor/test_shape.py

@@ -9,14 +9,11 @@
 
 import numpy as np
 import pytest
-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,
@@ -269,156 +266,6 @@ def test_concat_scalar():
     xr_assert_allclose(res, expected_res)
 
 
-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
-    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,)
-    fn = xr_function([x], y)
-    xr_assert_allclose(fn(xr_arange_like(x)), xr_arange_like(x).expand_dims("batch"))
-
-    # 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)
-
-    # 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")
-    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)
-
-    # Expand with multiple dims
-    x = xtensor("x", dims=(), shape=())
-    y = expand_dims(expand_dims(x, "a"), "b")
-    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)
-
-
-def test_expand_dims_implicit():
-    """Test expand_dims with default or symbolic sizes and dim=None."""
-
-    # 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")
-    expected = xr_arange_like(x).expand_dims("batch")
-    xr_assert_allclose(fn(xr_arange_like(x), 1), expected)
-
-    # Symbolic size > 1 (but expand only adds dim=1)
-    size_sym_4 = scalar("size_sym_4", dtype="int64")
-    y = expand_dims(x, "batch", size=size_sym_4)
-    fn = xr_function([x, size_sym_4], y, on_unused_input="ignore")
-    xr_assert_allclose(fn(xr_arange_like(x), 4), expected)
-
-    # Reversibility: expand then squeeze
-    x = xtensor("x", dims=("a", "b"), shape=(2, 3))
-    y = expand_dims(x, "batch")
-    z = squeeze(y, "batch")
-    fn = xr_function([x], z)
-    xr_assert_allclose(fn(xr_arange_like(x)), xr_arange_like(x))
-
-    # expand_dims with dim=None = no-op
-    x = xtensor("x", dims=("a",), shape=(3,))
-    y = expand_dims(x, None)
-    fn = xr_function([x], y)
-    xr_assert_allclose(fn(xr_arange_like(x)), xr_arange_like(x))
-
-    # broadcast after symbolic size
-    size_sym = scalar("size_sym", dtype="int64")
-    x = xtensor("x", dims=("a", "b"), shape=(2, 3))
-    y = expand_dims(x, "batch", size=size_sym)
-    z = y + y  # triggers shape alignment
-    fn = xr_function([x, size_sym], z, on_unused_input="ignore")
-    x_test = xr_arange_like(x)
-    out = fn(x_test, 1)
-    expected = x_test.expand_dims("batch") + x_test.expand_dims("batch")
-    xr_assert_allclose(out, expected)
-
-
-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")
-    with pytest.raises(ValueError, match="already exists"):
-        expand_dims(y, "city")
-
-    # Size = 0 is invalid
-    with pytest.raises(ValueError, match="size must be.*positive"):
-        expand_dims(x, "batch", size=0)
-
-    # Invalid dim type
-    with pytest.raises(TypeError):
-        expand_dims(x, 123)
-
-    # Invalid size type
-    with pytest.raises(TypeError):
-        expand_dims(x, "new", size=[1])
-
-    # Duplicate dimension creation
-    y = expand_dims(x, "new")
-    with pytest.raises(ValueError):
-        expand_dims(y, "new")
-
-    # Symbolic size with invalid runtime value
-    size_sym = scalar("size_sym", dtype="int64")
-    y = expand_dims(x, "batch", size=size_sym)
-    fn = xr_function([x, size_sym], y, on_unused_input="ignore")
-    with pytest.raises(Exception):
-        fn(xr_arange_like(x), 0)
-
-
 def test_squeeze_explicit_dims():
     """Test squeeze with explicit dimension(s)."""
 
@@ -487,13 +334,13 @@ def test_squeeze_implicit_dims():
     fn4 = xr_function([x4], y4)
     xr_assert_allclose(fn4(x4_test), x4_test.squeeze("b"))
 
-    # Reversibility with expand_dims
-    x5 = xtensor("x5", dims=("batch", "time", "feature"), shape=(2, 1, 3))
-    y5 = squeeze(x5, "time")
-    z5 = expand_dims(y5, "time")
-    fn5 = xr_function([x5], z5)
-    x5_test = xr_arange_like(x5)
-    xr_assert_allclose(fn5(x5_test).transpose(*x5_test.dims), x5_test)
+    # Reversibility with expand_dims (restore when expand_dims is implemented)
+    # x5 = xtensor("x5", dims=("batch", "time", "feature"), shape=(2, 1, 3))
+    # y5 = squeeze(x5, "time")
+    # z5 = expand_dims(y5, "time")
+    # fn5 = xr_function([x5], z5)
+    # x5_test = xr_arange_like(x5)
+    # xr_assert_allclose(fn5(x5_test).transpose(*x5_test.dims), x5_test)
 
     """
     This test documents that we intentionally don't squeeze dimensions with symbolic shapes