Fixing ExpandDims rewrite

Author: AllenDowney

pytensor/xtensor/rewriting/shape.py

@@ -1,17 +1,12 @@
 from pytensor.graph import node_rewriter
-from pytensor.raise_op import Assert
 from pytensor.tensor import (
     broadcast_to,
-    get_scalar_constant_value,
-    gt,
+    expand_dims,
     join,
     moveaxis,
     specify_shape,
     squeeze,
 )
-from pytensor.tensor import (
-    shape as tensor_shape,
-)
 from pytensor.xtensor.basic import tensor_from_xtensor, xtensor_from_tensor
 from pytensor.xtensor.rewriting.basic import register_lower_xtensor
 from pytensor.xtensor.shape import (
@@ -144,27 +139,30 @@ def local_squeeze_reshape(fgraph, node):
 @register_lower_xtensor
 @node_rewriter([ExpandDims])
 def local_expand_dims_reshape(fgraph, node):
-    """Rewrite ExpandDims to tensor.expand_dims and optionally broadcast_to or specify shape."""
+    """Rewrite ExpandDims using tensor operations."""
     x, size = node.inputs
     out = node.outputs[0]
-    # Lower to tensor.expand_dims(x, axis=0)
-    from pytensor.tensor import expand_dims as tensor_expand_dims
-
-    expanded = tensor_expand_dims(tensor_from_xtensor(x), 0)
-    # Optionally broadcast to the correct shape if size is not 1
-    from pytensor.tensor import broadcast_to
-
-    # Ensure size is positive
-    expanded = Assert(msg="size must be positive")(expanded, gt(size, 0))
-    # If size is not 1, broadcast
-    try:
-        static_size = get_scalar_constant_value(size)
-    except Exception:
-        static_size = None
-    if static_size is not None and static_size == 1:
-        result = expanded
+
+    # Convert inputs to tensors
+    x_tensor = tensor_from_xtensor(x)
+    size_tensor = tensor_from_xtensor(size)
+
+    # Get the new dimension name and position
+    new_axis = 0  # Always insert at front
+
+    # Use tensor operations
+    if out.type.shape[0] == 1:
+        # Simple case: just expand with size 1
+        result_tensor = expand_dims(x_tensor, new_axis)
     else:
-        # Broadcast to (size, ...)
-        new_shape = (size,) + tuple(tensor_shape(expanded))[1:]
-        result = broadcast_to(expanded, new_shape)
-    return [xtensor_from_tensor(result, out.type.dims)]
+        # 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))
+
+    # Preserve static shape information
+    result_tensor = specify_shape(result_tensor, out.type.shape)
+
+    # Convert result back to xtensor
+    result = xtensor_from_tensor(result_tensor, dims=out.type.dims)
+    return [result]

pytensor/xtensor/shape.py

@@ -6,6 +6,7 @@
 import numpy as np
 
 from pytensor.graph import Apply
+from pytensor.graph.basic import Constant
 from pytensor.scalar import discrete_dtypes, upcast
 from pytensor.tensor import as_tensor, get_scalar_constant_value
 from pytensor.tensor.exceptions import NotScalarConstantError
@@ -391,43 +392,47 @@ class ExpandDims(XOp):
     __props__ = ("dim",)
 
     def __init__(self, dim):
+        if not isinstance(dim, str):
+            raise TypeError(f"`dim` must be a string, got: {type(self.dim)}")
+
         self.dim = dim
 
     def make_node(self, x, size):
         x = as_xtensor(x)
 
-        if not isinstance(self.dim, str):
-            raise TypeError(f"`dim` must be a string or None, got: {type(self.dim)}")
-
         if self.dim in x.type.dims:
             raise ValueError(f"Dimension {self.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
+
+        # Check if size is a valid type before converting
+        if not (
+            isinstance(size, int | np.integer)
+            or (hasattr(size, "ndim") and getattr(size, "ndim", None) == 0)
         ):
             raise TypeError(
                 f"size must be an int or scalar variable, got: {type(size)}"
             )
 
-        # Convert size to tensor
-        size = as_tensor(size, ndim=0)
-
-        # Insert new dim at front
-        new_dims = (self.dim, *x.type.dims)
-
         # Determine shape
         try:
             static_size = get_scalar_constant_value(size)
         except NotScalarConstantError:
             static_size = None
+
         if static_size is not None:
             new_shape = (int(static_size), *x.type.shape)
         else:
             new_shape = (None, *x.type.shape)  # symbolic size
 
+        # Convert size to tensor
+        size = as_xtensor(size, dims=())
+
+        # Check if size is a constant and validate it
+        if isinstance(size, Constant) and size.data < 0:
+            raise ValueError(f"size must be 0 or positive, got: {size.data}")
+
+        # Insert new dim at front
+        new_dims = (self.dim, *x.type.dims)
+
         out = xtensor(
             dtype=x.type.dtype,
             shape=new_shape,

tests/xtensor/test_shape.py

@@ -540,10 +540,6 @@ def test_expand_dims_errors():
     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, match="Invalid type for `dim`"):
         expand_dims(x, 123)
@@ -557,13 +553,6 @@ def test_expand_dims_errors():
     with pytest.raises(ValueError, match="already exists"):
         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_expand_dims_multiple():
     """Test expanding multiple dimensions at once using a list of strings."""