From 4a7aed99a8b052e231746429b04363e532ed6c17 Mon Sep 17 00:00:00 2001 From: Chen Yufan Date: Mon, 31 Aug 2026 20:54:23 +0800 Subject: [PATCH] [Fix][Relax][Frontend][Torch] Validate flatten dims in `from_fx` Fixes: #20227 `_flatten_impl` normalized negative flatten dims but never checked their range or their ordering, so `from_fx` crashed with an internal `TypeError: reduce() of empty iterable with no initial value` on a traced model whose `start_dim` came after its `end_dim`. Out-of-range dims leaked an `IndexError`, and an out-of-range negative `start_dim` silently computed a wrong shape that only failed later inside `relax.op.reshape`. Normalize both dims against `max(rank, 1)`, validate each against `[-r, r-1]`, and reject `start_dim > end_dim` with a message matching torch's own. A 0-d input is now handled explicitly, matching `torch.flatten` on a scalar. Co-authored-by: Claude --- .../torch/base_fx_graph_translator.py | 34 +++++++++++-- tests/python/relax/test_frontend_from_fx.py | 50 +++++++++++++++++++ 2 files changed, 81 insertions(+), 3 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 d600987cdd7b..84280c196016 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -1938,13 +1938,41 @@ def _expand_as(self, node: fx.Node) -> relax.Var: def _flatten_impl(self, x, start_dim, end_dim) -> relax.Var: shape = self.shape_of(x) - start_dim = start_dim if start_dim >= 0 else len(shape) + start_dim - end_dim = end_dim if end_dim >= 0 else len(shape) + end_dim + rank = len(shape) + + # torch.flatten() normalizes its dims against a rank of at least one, so a 0-d + # input still accepts a start_dim/end_dim of 0 or -1. + dim_post_expr = max(rank, 1) + norm_start_dim = start_dim + dim_post_expr if start_dim < 0 else start_dim + norm_end_dim = end_dim + dim_post_expr if end_dim < 0 else end_dim + + # torch rejects invalid flatten dims only when the model is executed. fx.symbolic_trace + # does not execute it, so an invalid flatten reaches this converter as a traceable node + # and has to be rejected here instead of failing later on an empty reduce(). + if not 0 <= norm_start_dim < dim_post_expr: + raise ValueError( + f"flatten start_dim {start_dim} is out of range " + f"[-{dim_post_expr}, {dim_post_expr - 1}] for an input of rank {rank}" + ) + if not 0 <= norm_end_dim < dim_post_expr: + raise ValueError( + f"flatten end_dim {end_dim} is out of range " + f"[-{dim_post_expr}, {dim_post_expr - 1}] for an input of rank {rank}" + ) + if norm_start_dim > norm_end_dim: + raise ValueError("flatten() has invalid args: start_dim cannot come after end_dim") + + start_dim, end_dim = norm_start_dim, norm_end_dim + + # torch.flatten() on a 0-d input returns a 1-d tensor holding the single element. + if rank == 0: + return self.block_builder.emit(relax.op.reshape(x, [1])) + flattened = reduce(lambda x, y: x * y, [shape[i] for i in range(start_dim, end_dim + 1)]) new_shape = ( [shape[i] for i in range(0, start_dim)] + [flattened] - + [shape[i] for i in range(end_dim + 1, len(shape))] + + [shape[i] for i in range(end_dim + 1, rank)] ) return self.block_builder.emit(relax.op.reshape(x, new_shape)) diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py index a489977958c7..b56c11155161 100644 --- a/tests/python/relax/test_frontend_from_fx.py +++ b/tests/python/relax/test_frontend_from_fx.py @@ -1722,6 +1722,56 @@ def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor( verify_model(torch.nn.Flatten(2, -1), input_info, {}, expected1) +def test_flatten_invalid_dims(): + input_info = [([2, 3, 4], "float32")] + + class FlattenFunc(Module): + def forward(self, input): + return torch.flatten(input, 2, 1) + + class FlattenModule(Module): + def __init__(self): + super().__init__() + self.f = torch.nn.Flatten(2, 1) + + def forward(self, input): + return self.f(input) + + class FlattenOutOfRange(Module): + def forward(self, input): + return torch.flatten(input, 0, 3) + + # torch rejects these dims only when the model runs, and fx.symbolic_trace does not run + # it, so the invalid flatten reaches the frontend and has to be rejected there. + for model in (FlattenFunc(), FlattenModule()): + with pytest.raises(ValueError, match="start_dim cannot come after end_dim"): + from_fx(fx.symbolic_trace(model), input_info) + + with pytest.raises(ValueError, match="flatten end_dim 3 is out of range"): + from_fx(fx.symbolic_trace(FlattenOutOfRange()), input_info) + + +def test_flatten_scalar_input(): + input_info = [([], "float32")] + + class Flatten(Module): + def forward(self, input): + return torch.flatten(input) + + @tvm.script.ir_module + class expected1: + @R.function + def main(input_1: R.Tensor((), dtype="float32")) -> R.Tensor((1,), dtype="float32"): + # block 0 + with R.dataflow(): + lv: R.Tensor((1,), dtype="float32") = R.reshape(input_1, (1,)) + gv: R.Tensor((1,), dtype="float32") = lv + R.output(gv) + return gv + + verify_model(Flatten(), input_info, {}, expected1) + + def test_batchnorm2d(): input_info = [([1, 3, 10, 10], "float32")]