Cleaning up squeeze

Author: AllenDowney

pytensor/xtensor/rewriting/shape.py

@@ -162,28 +162,18 @@ 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."""
-    if not isinstance(node.op, Squeeze):
-        return False
-
-    [x] = node.inputs
-    in_dims = x.type.dims
+    """Rewrite Squeeze to tensor.squeeze."""
+    x = node.inputs[0]
     dim = node.op.dims
 
-    # 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:
-        dims_to_remove = [dim] if isinstance(dim, str) else dim
-        axes_to_squeeze = tuple(in_dims.index(d) for d in dims_to_remove)
-
-    # Nothing to squeeze? Just return input unchanged
-    if not axes_to_squeeze:
         return [x]
 
     x_tensor = tensor_from_xtensor(x)
+    x_dims = x.type.dims
+    dims_to_remove = [dim] if isinstance(dim, str) else dim
+    axes_to_squeeze = tuple(x_dims.index(d) for d in dims_to_remove)
     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

@@ -397,19 +397,21 @@ class Squeeze(XOp):
 
     __props__ = ("dims",)
 
-    def __init__(self, dims):
-        self.dims = dims
+    def __init__(self, dim):
+        self.dims = tuple(sorted(set(dim)))
 
     def make_node(self, x):
         x = as_xtensor(x)
 
         # Validate that dims exist and are size-1 if statically known
         dims_to_remove = []
+        x_dims = x.type.dims
+        x_shape = x.type.shape
         for d in self.dims:
-            if d not in x.type.dims:
+            if d not in x_dims:
                 raise ValueError(f"Dimension {d} not found in {x.type.dims}")
-            idx = x.type.dims.index(d)
-            dim_size = x.type.shape[idx]
+            idx = x_dims.index(d)
+            dim_size = x_shape[idx]
             if dim_size is not None and dim_size != 1:
                 raise ValueError(f"Dimension {d} has static size {dim_size}, not 1")
             dims_to_remove.append(idx)
@@ -454,9 +456,6 @@ def squeeze(x, dim=None):
     else:
         dims = tuple(dim)
 
-    # Normalize: deduplicate and sort
-    dims = tuple(sorted(set(dims)))
-
     if not dims:
         return x  # no-op if nothing to squeeze
 

tests/xtensor/test_shape.py

@@ -448,7 +448,7 @@ def test_squeeze_explicit_dims():
     # Redundant dimensions
     y3c = squeeze(x3, ["b", "b"])
     fn3c = xr_function([x3], y3c)
-    xr_assert_allclose(fn3c(x3_test), x3_test.squeeze("b"))
+    xr_assert_allclose(fn3c(x3_test), x3_test.squeeze(["b", "b"]))
 
     # Empty list = no-op
     y3d = squeeze(x3, [])
@@ -495,6 +495,22 @@ def test_squeeze_implicit_dims():
     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
+    (static_shape=None) even when they are 1 at runtime, while xarray does squeeze them.
+    """
+    # Create a tensor with a symbolic dimension that will be 1 at runtime
+    x = xtensor("x", dims=("a", "b", "c"))  # shape unknown
+    y = squeeze(x)  # implicit dim=None should not squeeze symbolic dimensions
+    x_test = xr_arange_like(xtensor(dims=x.dims, shape=(2, 1, 3)))
+    fn = xr_function([x], y)
+    res = fn(x_test)
+
+    # Our implementation should not squeeze the symbolic dimension
+    assert "b" in res.dims
+    # While xarray would squeeze it
+    assert "b" not in x_test.squeeze().dims
+
 
 def test_squeeze_errors():
     """Test error cases for squeeze."""