diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py index b0682f1f4ceb..9d44bbf3f6ed 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/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)) diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py index a0ba7971db0a..4cb27b75a6ab 100644 --- a/tests/python/relax/test_frontend_from_fx.py +++ b/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")]