This is an automated email from the ASF dual-hosted git repository.
tlopex 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 b3f51d1f28 [Fix][Relax][Frontend][Torch] Validate flatten dims in
`from_fx` (#20245)
b3f51d1f28 is described below
commit b3f51d1f28e602376c04c7c3d57b4f58e170954e
Author: Chen Yufan <[email protected]>
AuthorDate: Tue Sep 8 11:37:29 2026 +0800
[Fix][Relax][Frontend][Torch] Validate flatten dims in `from_fx` (#20245)
Fixes: #20227
## Summary
`tvm.relax.frontend.torch.from_fx` crashed with an internal
`TypeError: reduce() of empty iterable with no initial value` when a
traced model
contained a `flatten` whose `start_dim` comes after its `end_dim`.
torch rejects such a `flatten` with a clear `RuntimeError`, but only
when the model is
executed. `fx.symbolic_trace` does not execute the model, so the invalid
node reaches
the frontend as a perfectly traceable graph and has to be rejected
there. This PR
validates the dims in `_flatten_impl`, mirroring what `from_onnx` now
does for
`Flatten` (#20145).
Only the `from_fx` / `TorchFXImporter` path is affected.
`from_exported_program` runs
`run_decompositions()` by default, which lowers
`aten.flatten.using_ints` to
`aten.view`, so `_flatten_impl` is never reached there.
## Root cause
`_flatten_impl` in
`python/tvm/relax/frontend/torch/base_fx_graph_translator.py`
normalized negative dims but never checked their range or their
ordering:
```python
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
flattened = reduce(lambda x, y: x * y, [shape[i] for i in range(start_dim,
end_dim + 1)])
```
For `start_dim=2, end_dim=1` the `range` is empty, so `functools.reduce`
is called over
an empty iterable with no initial value. Out-of-range dims (`flatten(x,
0, 3)` on a
rank-3 input) leaked an `IndexError` out of `shape[i]` instead.
Both entry points that reach this helper are affected: `torch.flatten`
(dispatched via
`_flatten`) and `torch.nn.Flatten` (via `_flatten_module`).
## Fix
Normalize both dims against `max(rank, 1)`, validate each against `[-r,
r-1]`, and
reject `start_dim > end_dim` with torch's own wording:
```python
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
...
if norm_start_dim > norm_end_dim:
raise ValueError("flatten() has invalid args: start_dim cannot come
after end_dim")
```
A 0-d input is handled explicitly: torch normalizes flatten dims against
a rank of at
least one, so `torch.flatten(scalar)` is valid and returns a 1-d tensor
holding the
single element. The old code hit the same empty-`reduce` crash for that
input.
## Validation
Built from source (CPU-only, `USE_LLVM=OFF`) on Linux, torch 2.13.0+cpu,
Python 3.14.
Behavior on the reported cases, measured before and after the change:
| Input shape | dims | Before | After |
|---|---|---|---|
| `(2,3,4)` | `(2,1)` | `TypeError: reduce() of empty iterable with no
initial value` | `ValueError: flatten() has invalid args: start_dim
cannot come after end_dim` |
| `(2,3,4)` | `(0,3)` | `IndexError: ShapeExpr index out of range` |
`ValueError: flatten end_dim 3 is out of range [-3, 2] for an input of
rank 3` |
| `(2,3,4)` | `(-4,2)` | `ValueError: Reshape expects the new shape to
be convertible from the old shape` | `ValueError: flatten start_dim -4
is out of range [-3, 2] for an input of rank 3` |
| `()` (0-d) | `(0,-1)` | `TypeError: reduce() of empty iterable with no
initial value` | converts, output `(1,)` |
The `(-4,2)` row is worth calling out: an out-of-range negative
`start_dim` did not fail in
`_flatten_impl` at all. It normalized to `-1`, so `range(-1, 3)`
silently folded the last
dimension into the product and emitted a `reshape` to `(96,)` for a
24-element tensor, which
only failed later inside `relax.op.reshape`. That is a mis-computation,
not just a crash.
Valid dims are unaffected — `(1,3,10,10)` with `(2,-1)`, and `(2,3,4)`
with `(1,2)`, `(0,-1)`,
`(-3,-1)` and `(2,2)`, all convert to the same shapes as before.
Regression run of the whole `from_fx` suite, base commit vs. this
branch:
| | base (`HEAD~1`) | this branch |
|---|---|---|
| passed | 164 | 166 (+2 new tests) |
| failed | 16 | 16 |
| skipped | 1 | 1 |
The failure sets are identical line for line — the 16 failures
(`test_dtypes[*]`, `test_round`)
are pre-existing on the base commit in this environment and unrelated to
this change.
Tests added to `tests/python/relax/test_frontend_from_fx.py`:
- `test_flatten_invalid_dims` — `start_dim > end_dim` through both
`torch.flatten` and
`torch.nn.Flatten`, plus an out-of-range `end_dim`
- `test_flatten_scalar_input` — 0-d input flattens to shape `(1,)`
```
$ pytest tests/python/relax/test_frontend_from_fx.py -k flatten -v
test_flatten PASSED
test_flatten_invalid_dims PASSED
test_flatten_scalar_input PASSED
3 passed, 180 deselected
$ pre-commit run --files
python/tvm/relax/frontend/torch/base_fx_graph_translator.py \
tests/python/relax/test_frontend_from_fx.py
ruff check ....... Passed
ruff format ...... Passed
(all hooks passed)
```
## Files changed
- `python/tvm/relax/frontend/torch/base_fx_graph_translator.py` —
`_flatten_impl`:
normalize and validate dims, handle a 0-d input.
- `tests/python/relax/test_frontend_from_fx.py` — regression tests.
## Note for reviewers
The 0-d handling (`dim_post_expr = max(rank, 1)` plus the `rank == 0`
branch) is
separable from the reported bug. It fixes the same empty-`reduce` crash
for scalar
inputs and keeps the new range check from rejecting a valid
`torch.flatten(scalar)`,
but if you would rather keep this PR to exactly the reported case I am
happy to drop
that hunk and its test.
---
This change was developed with AI assistance. The build, the
before/after comparison
and the test runs reported above were all executed locally against this
branch as
submitted.
Co-authored-by: Claude <[email protected]>
---
.../frontend/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 f5f8afd441..f6781f026f 100644
--- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
+++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
@@ -1943,13 +1943,41 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
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 0759aab1ee..7189b3cb24 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -1722,6 +1722,56 @@ def test_flatten():
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")]
