[Bug][Relax][Frontend][Torch] `from_fx` crashes with obscure internal `TypeError: reduce() of empty iterable` for `flatten` with `start_dim > end_dim`, instead of a clear "invalid flatten args" error

[siyiweigeHEW]

Expected behavior

A model whose flatten uses invalid dims (start_dim > end_dim) is rejected by PyTorch itself at execution time:

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:

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:

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

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