siyiweigeHEW opened a new issue, #20227:
URL: https://github.com/apache/tvm/issues/20227
### Expected behavior
A model whose `flatten` uses invalid dims (`start_dim > end_dim`) is
rejected by PyTorch itself at execution time:
```python
torch.flatten(x, 2, 1) # RuntimeError: flatten() has invalid args:
start_dim cannot come after end_dim
torch.nn.Flatten(2, 1)(x) # same
```
`torch.fx.symbolic_trace` does **not** execute the model and does **not**
validate `flatten` dims (the check only fires on execution), so such a graph is
perfectly traceable:
```python
import torch
gm = torch.fx.symbolic_trace(torch.nn.Flatten(2, 1)) # succeeds
```
When such a graph is passed to the TVM torch frontend, the conversion should
reject the invalid `flatten` with a clear error message (matching torch's
`RuntimeError` / `ValueError`), not crash with an internal implementation
detail.
### Actual behavior
`tvm.relax.frontend.torch.from_fx` (via `TorchFXImporter`) dispatches the
`flatten` node to `_flatten` → `_flatten_impl`
(`python/tvm/relax/frontend/torch/base_fx_graph_translator.py:1771-1787`),
which computes the flattened extent as:
```python
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` on a 3-D tensor, `range(2, 2)` is empty, so
`functools.reduce` is called over an empty iterable with no initial value and
raises:
```
TypeError: reduce() of empty iterable with no initial value
```
The same crash is reachable both from a raw `torch.flatten(x, 2, 1)` call in
the traced module and from `torch.nn.Flatten(start_dim=2, end_dim=1)`. A legal
`flatten(1, 2)` converts and runs correctly (`(2, 12)`).
### Environment
- OS: Linux
- TVM: v0.24.dev0 (main branch, commit `262c6d2e0`, built 2026-02-11)
- Python: 3.11
- torch: 2.10.0+cu128
### Steps to reproduce
```python
import torch
from tvm.relax.frontend.torch import from_fx
# Variant 1: raw torch.flatten with invalid dims
class BadFlatten(torch.nn.Module):
def forward(self, x):
return torch.flatten(x, 2, 1)
# Variant 2: nn.Flatten with start_dim > end_dim
class BadFlattenModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.flatten = torch.nn.Flatten(start_dim=2, end_dim=1)
def forward(self, x):
return self.flatten(x)
x = torch.randn(2, 3, 4)
# torch reference: both are rejected on execution
for name, m in [("torch.flatten(x,2,1)", BadFlatten()), ("nn.Flatten(2,1)",
BadFlattenModule())]:
try:
m(x)
print(f"torch native {name}: OK")
except Exception as e:
print(f"torch native {name}: {type(e).__name__}: {e}")
# TVM: symbolic_trace succeeds (no validation), from_fx crashes
for gm in (torch.fx.symbolic_trace(BadFlatten().eval()),
torch.fx.symbolic_trace(BadFlattenModule().eval())):
try:
from_fx(gm, input_info=[((2, 3, 4), "float32")])
print("from_fx: OK")
except Exception as e:
print(f"from_fx: {type(e).__name__}: {e}")
```
Actual output:
```
torch native torch.flatten(x,2,1): RuntimeError: flatten() has invalid args:
start_dim cannot come after end_dim
torch native nn.Flatten(2,1): RuntimeError: flatten() has invalid args:
start_dim cannot come after end_dim
from_fx: TypeError: reduce() of empty iterable with no initial value
from_fx: TypeError: reduce() of empty iterable with no initial value
```
Full traceback (single variant):
```
File "tvm/relax/frontend/torch/fx_translator.py", line 1256, in from_fx
return
TorchFXImporter(default_image_layout=default_image_layout).from_fx(
File "tvm/relax/frontend/torch/fx_translator.py", line 1134, in from_fx
self.env[node] = self.convert_map[func_name](node)
File "tvm/relax/frontend/torch/base_fx_graph_translator.py", line 1787, in
_flatten
return self._flatten_impl(x, start_dim, end_dim)
File "tvm/relax/frontend/torch/base_fx_graph_translator.py", line 1775, in
_flatten_impl
flattened = reduce(lambda x, y: x * y, [shape[i] for i in
range(start_dim, end_dim + 1)])
TypeError: reduce() of empty iterable with no initial value
```
### Notes / suggested fix
- Root cause: `_flatten_impl` does not validate that `start_dim <= end_dim`
after negative-index normalization. On an empty `range` it hits
`functools.reduce` with no initial value.
- Suggest validating the dims and raising a clear error (e.g. matching
torch's `"flatten() has invalid args: start_dim cannot come after end_dim"`),
or equivalently guarding the empty-range case.
- Scope: this affects the `from_fx` / `TorchFXImporter` entry point (where
`flatten` is dispatched to `_flatten`). The `from_exported_program` path is not
affected: the module-level `from_exported_program` runs
`exported_program.run_decompositions()` by default
(`python/tvm/relax/frontend/torch/exported_program_translator.py:1841-1842`),
which decomposes `aten.flatten.using_ints` into `aten.view` (shape computed at
decompose time), so `_flatten_impl` is never reached there.
### Triage
* needs-triage
* bug
* relax
* frontend/torch
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
