This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 7f619733d1 [Fix][Relax][Torch] Align retained expand dimensions by
trailing rank (#20137)
7f619733d1 is described below
commit 7f619733d130918e0bd63ea6b732663190539806
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Tue Aug 25 17:13:10 2026 -0700
[Fix][Relax][Torch] Align retained expand dimensions by trailing rank
(#20137)
`torch.Tensor.expand` aligns existing input dimensions with the trailing
dimensions of the requested shape when new leading dimensions are
introduced. The Torch frontend previously resolved `-1` using the output
dimension index directly. For an input shaped `(tokens, 3)`,
`x.expand(2, -1, -1)` thus selected the wrong input dimensions and
eventually accessed beyond the input rank.
This PR:
- Accounts for newly introduced leading dimensions when resolving `-1`
- Aligns retained dimensions against the trailing input rank
- Rejects `-1` for new leading dimensions, matching PyTorch semantics.=
- Preserves the FX metadata fallback when the Relax input shape is
unavailable
---
.../relax/frontend/torch/base_fx_graph_translator.py | 11 +++++++++--
.../relax/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 85a0d956ca..b0bb14ac95 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 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
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 b2aa982fdf..e9d2ac8b70 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -5529,6 +5529,26 @@ def
test_derived_input_dimension_without_exported_program_decomposition():
assert tvm.arith.Analyzer().can_prove_equal(y_shape[1], x_shape[1] * 2)
+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):
