[Fix][Relax][Torch] Align retained expand dimensions by trailing rank

python/tvm/relax/frontend/torch/base_fx_graph_translator.py

@@ -1905,14 +1905,21 @@ def _expand(self, node: fx.Node) -> relax.Var:
         x = args[0]
         broadcast_shape = []
         in_shape = self.shape_of(x)
+        input_rank = len(in_shape) if in_shape is not None else None
+        if input_rank is None and hasattr(node.args[0], "meta") and "val" in node.args[0].meta:
+            input_rank = len(node.args[0].meta["val"].shape)
+        rank_offset = len(sizes) - input_rank if input_rank is not None else 0
         for idx, i in enumerate(sizes):
             if isinstance(i, int) and i == -1:
+                input_idx = idx - rank_offset
+                if input_idx < 0:
+                    raise ValueError(f"Cannot use -1 in expand for new leading dim {idx}")
                 if in_shape is not None:
-                    broadcast_shape.append(in_shape[idx])
+                    broadcast_shape.append(in_shape[input_idx])
                 elif hasattr(node.args[0], "meta") and "val" in node.args[0].meta:
                     # Fallback: get shape from FX node metadata (FakeTensor)
                     fake_shape = node.args[0].meta["val"].shape
-                    broadcast_shape.append(fake_shape[idx])
+                    broadcast_shape.append(fake_shape[input_idx])
                 else:
                     raise ValueError(
                         f"Cannot use -1 in expand for dim {idx} when input shape is unknown"

tests/python/relax/test_frontend_from_exported_program.py

@@ -5507,6 +5507,26 @@ def main(x: R.Tensor((2, 8, 4), dtype="float32")) -> R.Tuple(
     verify_model(SliceStaticModel(), example_args_static, {}, ExpectedStatic)
 
 
+def test_expand_with_new_leading_dimension():
+    class ExpandLeading(torch.nn.Module):
+        def forward(self, x):
+            return x.expand(2, -1, -1)
+
+    tokens = torch.export.Dim("tokens", min=1, max=8)
+    exported_program = export(
+        ExpandLeading(),
+        args=(torch.randn(4, 3),),
+        dynamic_shapes={"x": {0: tokens}},
+    )
+    mod = from_exported_program(exported_program)
+
+    input_shape = mod["main"].params[0].ty.shape.values
+    output_shape = mod["main"].ret_ty.fields[0].shape.values
+    assert tvm.arith.Analyzer().can_prove_equal(output_shape[0], 2)
+    assert tvm.arith.Analyzer().can_prove_equal(output_shape[1], input_shape[0])
+    assert tvm.arith.Analyzer().can_prove_equal(output_shape[2], input_shape[1])
+
+
 def test_split():
     class Chunk(Module):
         def forward(self, input):