This is an automated email from the ASF dual-hosted git repository.
jikechao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 66394cd89f [Fix][Relax][Frontend][Torch] Validate `dim` in the
`squeeze` converter (#20322)
66394cd89f is described below
commit 66394cd89ff5a74c633a8a8ffada00ecf065eb29
Author: HuEnwei <[email protected]>
AuthorDate: Thu Oct 15 08:18:36 2026 +0800
[Fix][Relax][Frontend][Torch] Validate `dim` in the `squeeze` converter
(#20322)
Fixes #20321.
Reject out-of-range squeeze dimensions with a clear ValueError instead of
silently dropping axes or raising an internal error. Add regression tests for
invalid dimensions and valid tuple inputs.
Also fix existing clang-format violations in include/tvm/tirx/type.h and
src/ir/expr.cc so that whole-tree lint checks pass.
---
.../frontend/torch/base_fx_graph_translator.py | 22 ++++++-------
tests/python/relax/test_frontend_from_fx.py | 36 ++++++++++++++++++++++
2 files changed, 46 insertions(+), 12 deletions(-)
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 3706411e43..b6fc656e98 100644
--- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
+++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
@@ -2621,18 +2621,16 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
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
+ # fx.symbolic_trace never executes the model, so torch's IndexError
for an
+ # out-of-range dim does not fire; reject the axis here instead of
dropping it.
+ 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 53e33ce334..bf6f04cf54 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -2565,6 +2565,42 @@ def test_squeeze():
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)
+
+ # fx.symbolic_trace never runs the model, so torch's IndexError does not
fire here.
+ 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")]
