[Fix][Relax][Frontend][Torch] Validate flatten dims in from_fx

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))
 

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")]