From 3f634a1172d1252f369ba403423ec6b0730044bd Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Wed, 12 Aug 2026 03:01:00 -0400 Subject: [PATCH] [FIX][RELAX][Torch] Align retained expand dimensions by trailing rank --- .../torch/base_fx_graph_translator.py | 11 ++++++++-- .../test_frontend_from_exported_program.py | 20 +++++++++++++++++++ 2 files changed, 29 insertions(+), 2 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 85a0d956ca1d..b0bb14ac95ad 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/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" diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 3799b8ed95f5..46f8eb5e77b8 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/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):