[Fix][Relax][Frontend][Torch] Validate dim in the squeeze converter

python/tvm/relax/frontend/torch/base_fx_graph_translator.py

@@ -2527,18 +2527,19 @@ def _squeeze(self, node: fx.Node) -> relax.Var:
         if dim is None:
             dim = node.kwargs.get("dims", None)
 
-        # If dims is a list, filter out axes where dimension is not 1
-        # This is needed because PyTorch decomposition may pass all axes
-        if isinstance(dim, list | tuple) and len(dim) > 0:
-            shape = self.shape_of(x)
-            # Filter to only include axes where the dimension is 1
-            valid_dims = []
-            for d in dim:
-                axis = d if d >= 0 else len(shape) + d
-                if axis < len(shape):
-                    valid_dims.append(d)
-            # If no valid dims, use None to squeeze all size-1 dimensions
-            dim = valid_dims if valid_dims else None
+        # torch rejects an out-of-range dim with IndexError, but only when the model is
+        # executed. fx.symbolic_trace does not execute it, so the invalid axis reaches this
+        # converter and has to be rejected here. An out-of-range axis in a list used to be
+        # filtered out, and a list whose axes were all out of range fell back to
+        # `dim=None` -- silently squeezing every size-1 dim instead of reporting the axis.
+        if dim is not None:
+            rank = len(self.shape_of(x))
+            for d in dim if isinstance(dim, list | tuple) else [dim]:
+                if isinstance(d, int) and not -rank <= d < rank:
+                    raise ValueError(
+                        f"squeeze dim {d} is out of range "
+                        f"[-{rank}, {rank - 1}] for an input of rank {rank}"
+                    )
 
         return self.block_builder.emit(relax.op.squeeze(x, dim))
 

tests/python/relax/test_frontend_from_fx.py

@@ -2721,6 +2721,46 @@ def main(inp_0: R.Tensor((3, 1, 4, 1), dtype="float32")) -> R.Tensor(
     verify_model(Squeeze2(), input_info, {}, Expected2)
 
 
+def test_squeeze_out_of_range_dim():
+    input_info = [([1, 2, 1], "float32")]
+
+    class Squeeze(Module):
+        def __init__(self, dim):
+            super().__init__()
+            self.dim = dim
+
+        def forward(self, input):
+            return input.squeeze(self.dim)
+
+    # torch rejects an out-of-range dim with IndexError, but only when the model is
+    # executed, and fx.symbolic_trace does not execute it, so the invalid axis reaches the
+    # frontend and has to be rejected there. Note that a tuple whose axes are all out of
+    # range used to be filtered down to an empty list and silently reinterpreted as
+    # squeeze(None), removing every size-1 dim instead of reporting the bad axis.
+    for dim in ((5,), 3, -4, (-4,), (0, 5)):
+        with pytest.raises(ValueError, match="squeeze dim .* is out of range"):
+            from_fx(fx.symbolic_trace(Squeeze(dim)), input_info)
+
+    # In-range dims, including the tuple form, still convert.
+    class SqueezeTuple(Module):
+        def forward(self, input):
+            return input.squeeze((0, 2))
+
+    @tvm.script.ir_module
+    class Expected:
+        @R.function
+        def main(
+            inp_0: R.Tensor((1, 2, 1), dtype="float32"),
+        ) -> R.Tensor((2,), dtype="float32"):
+            with R.dataflow():
+                lv: R.Tensor((2,), dtype="float32") = R.squeeze(inp_0, axis=[0, 2])
+                gv: R.Tensor((2,), dtype="float32") = lv
+                R.output(gv)
+            return gv
+
+    verify_model(SqueezeTuple(), input_info, {}, Expected)
+
+
 def test_unsqueeze():
     input_info = [([1, 3, 10, 10], "float32")]