squeeze.dims: keep listed dims whose size is not 1

coreai_torch/_aten_to_core.py

@@ -2737,11 +2737,19 @@ def replace_softmax(
 def replace_squeeze_dims(
     values_map: dict[str, Value], node: fx.Node, loc: Location
 ) -> Value:
-    """Converts aten.squeeze.dims to coreai.shrink_dims."""
+    """Converts aten.squeeze.dims to coreai.shrink_dims.
+
+    PyTorch keeps a listed dim whose size is not 1 (x.squeeze() lists every
+    dim), so only the dims of static size 1 in the exported input are shrunk.
+    """
     x = _get_operand(values_map, node, 0)
-    return coreai.shrink_dims(
-        x, [d + x.type.rank if d < 0 else d for d in node.args[1]]
-    )
+    in_shape = node.args[0].meta["val"].shape
+    dims = [
+        d + x.type.rank if d < 0 else d
+        for d in node.args[1]
+        if isinstance(in_shape[d], int) and in_shape[d] == 1
+    ]
+    return coreai.shrink_dims(x, dims) if dims else x
 
 
 def replace_sum_dim_intlist(

tests/ops/test_ops.py

@@ -5768,6 +5768,28 @@ def forward(self, x: Tensor) -> Tensor:
     await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes)
 
 
+@pytest.mark.parametrize("dynamic", [False, True])
+@pytest.mark.parametrize(
+    "dims",
+    [
+        [1],  # size 4: PyTorch returns the input unchanged
+        [0, 1],  # only dim 0 is removed
+        None,  # x.squeeze(), exported as squeeze.dims over every dim
+    ],
+)
+async def test_squeeze_dims_not_size_one(dims: list[int] | None, dynamic: bool) -> None:
+    """A listed dim whose size is not 1 stays in the output, as in PyTorch."""
+    x = torch.rand(1, 4, 8)
+
+    class SqueezeModel(nn.Module):
+        def forward(self, x: Tensor) -> Tensor:
+            return x.squeeze() if dims is None else torch.squeeze(x, dims)
+
+    model = SqueezeModel().eval()
+    dynamic_shapes = make_dynamic_shapes(x=[None, "x1", "x2"]) if dynamic else None
+    await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes)
+
+
 class TestIndexTensor:
     """Test suite for index.Tensor operation."""