Updating squeeze

Author: AllenDowney

pytensor/xtensor/rewriting/shape.py

@@ -154,38 +154,28 @@ def local_expand_dims_reshape(fgraph, node):
 @register_xcanonicalize
 @node_rewriter([Squeeze])
 def local_squeeze_reshape(fgraph, node):
-    """Rewrite rule to convert squeeze to pytensor.tensor.squeeze."""
+    """Rewrite rule to convert Squeeze to pytensor.tensor.squeeze."""
     if not isinstance(node.op, Squeeze):
         return False
 
-    x = node.inputs[0]
+    [x] = node.inputs
+    in_dims = x.type.dims
     dim = node.op.dim
 
-    # Convert single dimension to iterable for consistent handling
-    dims_to_remove = [dim] if isinstance(dim, str) else dim
-
-    if dims_to_remove is not None:
-        # Validate dimensions exist and have size 1
-        dim_indices = []
-        for d in dims_to_remove:
-            if d not in x.type.dims:
-                return False
-            dim_idx = x.type.dims.index(d)
-            # Only check shape != 1 if the shape is not None (symbolic)
-            if x.type.shape[dim_idx] is not None and x.type.shape[dim_idx] != 1:
-                return False
-            dim_indices.append(dim_idx)
+    # Determine which axes to squeeze
+    if dim is None:
+        # Infer axes by comparing input and output dims
+        out_dims = node.outputs[0].type.dims
+        axes_to_squeeze = tuple(i for i, d in enumerate(in_dims) if d not in out_dims)
     else:
-        # Find all dimensions of size 1
-        dim_indices = [i for i, s in enumerate(x.type.shape) if s == 1]
-        if not dim_indices:
-            return False
+        dims_to_remove = [dim] if isinstance(dim, str) else dim
+        axes_to_squeeze = tuple(in_dims.index(d) for d in dims_to_remove)
 
-    # Create new dimensions list
-    new_dims = [d for i, d in enumerate(x.type.dims) if i not in dim_indices]
+    # Nothing to squeeze? Just return input unchanged
+    if not axes_to_squeeze:
+        return [x]
 
-    # Convert to tensor and use pytensor.tensor.squeeze
     x_tensor = tensor_from_xtensor(x)
-    x_tensor_squeezed = squeeze(x_tensor, axis=tuple(dim_indices))
-    new_out = xtensor_from_tensor(x_tensor_squeezed, dims=tuple(new_dims))
+    x_tensor_squeezed = squeeze(x_tensor, axis=axes_to_squeeze)
+    new_out = xtensor_from_tensor(x_tensor_squeezed, dims=node.outputs[0].type.dims)
     return [new_out]

pytensor/xtensor/shape.py

@@ -364,45 +364,62 @@ class Squeeze(XOp):
 
     Parameters
     ----------
-    dim : str or None or iterable of str
-        The name(s) of the dimension(s) to remove. If None, all dimensions of size 1 will be removed.
+    dim : str, None, or iterable of str
+        The name(s) of the dimension(s) to remove. If None, all dimensions
+        that are statically known to have size 1 will be removed.
+        Dimensions with symbolic shape will not be removed unless explicitly named.
+
+        Note: Unlike NumPy/xarray, if dim is None, only dimensions known to
+        be size 1 at graph construction time will be removed, even if they happen
+        to be size 1 at runtime.
     """
 
+    __props__ = ("dim",)
+
     def __init__(self, dim=None):
-        self.dim = dim
+        if dim is None:
+            self.dim = None
+        else:
+            dims = [dim] if isinstance(dim, str) else dim
+            if not all(isinstance(d, str) for d in dims):
+                raise TypeError(f"All dimension names must be strings: got {dims}")
+            # Deduplicate and sort to make __props__ deterministic and hashable
+            self.dim = tuple(sorted(set(dims)))
+
+            if not self.dim:
+                warnings.warn(
+                    "Squeeze received an empty dim list — no dimensions will be removed."
+                )
 
     def make_node(self, x):
         x = as_xtensor(x)
 
-        # Convert single dimension to iterable for consistent handling
-        dims_to_remove = [self.dim] if isinstance(self.dim, str) else self.dim
+        if self.dim is None:
+            # Auto-detect static size-1 dimensions
+            dims_to_remove = [d for d, s in zip(x.type.dims, x.type.shape) if s == 1]
+            if not dims_to_remove:
+                raise ValueError("No dimensions of size 1 to remove")
+        else:
+            dims_to_remove = list(self.dim)
 
-        if dims_to_remove is not None:
-            # Validate dimensions exist and have size 1
+            # Validate existence and static shape (when possible)
             for dim in dims_to_remove:
                 if dim not in x.type.dims:
                     raise ValueError(f"Dimension {dim} not found")
                 dim_idx = x.type.dims.index(dim)
-                # Only raise an error if the shape is statically known and not 1.
-                # If the shape is None (symbolic), defer the error to runtime.
-                if x.type.shape[dim_idx] is not None and x.type.shape[dim_idx] != 1:
-                    raise ValueError(
-                        f"Dimension {dim} has size {x.type.shape[dim_idx]}, not 1"
-                    )
-            # Get indices of dimensions to remove
-            dim_indices = [x.type.dims.index(dim) for dim in dims_to_remove]
-        else:
-            # Find all dimensions of size 1
-            dim_indices = [i for i, s in enumerate(x.type.shape) if s == 1]
-            if not dim_indices:
-                raise ValueError("No dimensions of size 1 to remove")
+                shape = x.type.shape[dim_idx]
+                if shape is not None and shape != 1:
+                    raise ValueError(f"Dimension {dim} has size {shape}, not 1")
+
+        dim_indices = [x.type.dims.index(dim) for dim in dims_to_remove]
 
-        # Create new dimensions and shape lists
         new_dims = [d for i, d in enumerate(x.type.dims) if i not in dim_indices]
         new_shape = [s for i, s in enumerate(x.type.shape) if i not in dim_indices]
 
         output = xtensor(
-            dtype=x.type.dtype, shape=tuple(new_shape), dims=tuple(new_dims)
+            dtype=x.type.dtype,
+            shape=tuple(new_shape),
+            dims=tuple(new_dims),
         )
         return Apply(self, [x], [output])
 

tests/xtensor/test_shape.py

@@ -386,3 +386,71 @@ def test_squeeze():
     x3d = xtensor("x3d", dims=("row", "col", "batch"), shape=(2, 3, 4))
     with pytest.raises(ValueError):
         squeeze(x3d)
+
+
+def test_squeeze_additional_cases():
+    # Redundant dimensions: squeeze(["b", "b"]) should behave like squeeze(["b"])
+    x1 = xtensor("x1", dims=("a", "b", "c"), shape=(2, 1, 1))
+    y1 = squeeze(x1, ["b", "b"])
+    fn1 = xr_function([x1], y1)
+    x1_test = xr_arange_like(x1)
+    expected1 = x1_test.squeeze(["b"])
+    xr_assert_allclose(fn1(x1_test), expected1)
+
+    # Symbolic shape: dim is 1 at runtime → should squeeze successfully
+    x2 = xtensor("x2", dims=("a", "b", "c"))  # shape unknown
+    y2 = squeeze(x2, "b")
+    fn2 = xr_function([x2], y2)
+    x2_test = xr_arange_like(xtensor(dims=x2.dims, shape=(2, 1, 3)))
+    expected2 = x2_test.squeeze("b")
+    xr_assert_allclose(fn2(x2_test), expected2)
+
+    # Symbolic shape: dim is not 1 at runtime → should raise
+    x3 = xtensor("x3", dims=("a", "b", "c"))  # shape unknown
+    y3 = squeeze(x3, "b")
+    fn3 = xr_function([x3], y3)
+    x3_test = xr_arange_like(xtensor(dims=x3.dims, shape=(2, 2, 3)))
+    with pytest.raises(Exception):
+        fn3(x3_test)
+
+    # Reversibility: squeeze then expand_dims should restore original
+    # TODO: uncomment when we have expand_dims
+    # x4 = xtensor("x4", dims=("batch", "time", "feature"), shape=(2, 1, 3))
+    # y4 = squeeze(x4, "time")
+    # z4 = expand_dims(y4, "time")
+    # fn4 = xr_function([x4], z4)
+    # x4_test = xr_arange_like(x4)
+    # xr_assert_allclose(fn4(x4_test), x4_test)
+
+
+def test_squeeze_extra_cases():
+    # 1. Order of dims shouldn't affect result
+    x1 = xtensor("x1", dims=("a", "b", "c"), shape=(2, 1, 1))
+    y1 = squeeze(x1, ["b", "c"])
+    y2 = squeeze(x1, ["c", "b"])
+    fn1 = xr_function([x1], y1)
+    fn2 = xr_function([x1], y2)
+    x1_test = xr_arange_like(x1)
+    xr_assert_allclose(fn1(x1_test), fn2(x1_test))
+
+    # 2. Empty list of dims = no-op
+    x2 = xtensor("x2", dims=("a", "b", "c"), shape=(2, 1, 1))
+    y2 = squeeze(x2, [])
+    fn2 = xr_function([x2], y2)
+    x2_test = xr_arange_like(x2)
+    xr_assert_allclose(fn2(x2_test), x2_test)
+
+    # 3. Explicit squeeze of all size-1 dims via dim=None
+    x3 = xtensor("x3", dims=("a", "b"), shape=(1, 1))
+    y3 = squeeze(x3)
+    fn3 = xr_function([x3], y3)
+    x3_test = xr_arange_like(x3)
+    xr_assert_allclose(fn3(x3_test), x3_test.squeeze())
+
+    # 4. Static + symbolic shape mix: squeeze symbolic 1-sized dim
+    x4 = xtensor("x4", dims=("a", "b", "c"), shape=(None, 1, 3))
+    y4 = squeeze(x4, "b")
+    x4_test = xr_arange_like(xtensor(dims=x4.dims, shape=(4, 1, 3)))
+    fn4 = xr_function([x4], y4)
+    expected4 = x4_test.squeeze("b")
+    xr_assert_allclose(fn4(x4_test), expected4)