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 af4c3f4d50 [Fix][Relax][ONNX] Preserve rank-expanding Expand (#19992)
af4c3f4d50 is described below
commit af4c3f4d505d27aae8cdf208e2304d67d20ac2cf
Author: Vic Wen <[email protected]>
AuthorDate: Tue Jul 14 02:42:15 2026 +0800
[Fix][Relax][ONNX] Preserve rank-expanding Expand (#19992)
### What changed
Record the original input rank before `Expand` left-pads the input shape
for broadcast validation. The no-op fast path now returns the input
unchanged only when both the padded shape and the original rank match
the target.
A regression test covers expanding `[1]` to `[1, 1]` when the target is
represented as a Relax `ShapeExpr`.
### Why
ONNX `Expand` right-aligns dimensions and may increase tensor rank by
adding leading dimensions. Previously, a rank-expanding broadcast could
look like a no-op after the frontend padded the input shape, causing it
to return the original lower-rank tensor. Downstream operators could
then receive inconsistent ranks.
This fixes the focused bug tracked in #19991 and is part of the
investigation and fixes for #19971. It does not close #19971 because the
attached model exposes additional independent importer issues after this
`Concat` failure is resolved.
Fixes #19991
Part of #19971
### Validation
- `python -m pytest
tests/python/relax/test_frontend_onnx.py::test_expand -q`
- A/B checked the model attached to #19971: the base revision reproduces
`Concat expects all input tensors to have same ndim`, while this change
advances beyond that `Concat`.
Signed-off-by: viiccwen <[email protected]>
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 3 ++-
tests/python/relax/test_frontend_onnx.py | 19 +++++++++++++++++++
2 files changed, 21 insertions(+), 1 deletion(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index d340c696b4..70e90d3731 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -2875,6 +2875,7 @@ class Expand(OnnxOpConverter):
shape = inputs[1]
if isinstance(shape, relax.ShapeExpr):
data_shape = list(data.ty.shape)
+ data_ndim = len(data_shape)
target_shape = list(shape.values)
original_data_shape = [
dim.value if hasattr(dim, "value") else str(dim) for dim in
data_shape
@@ -2918,7 +2919,7 @@ class Expand(OnnxOpConverter):
f"the same value or one of them to be 1."
)
# For dynamic shapes, let broadcast_to handle it
- if target_shape == data_shape:
+ if target_shape == data_shape and len(target_shape) == data_ndim:
return data
return relax.op.broadcast_to(data, relax.ShapeExpr(target_shape))
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 126c909983..c1bf8802f0 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -5601,6 +5601,19 @@ def test_expand():
R.output(gv)
return gv
+ @I.ir_module
+ class ExpectedHigherRankSamePaddedShape:
+ @R.function
+ def main(
+ in_: R.Tensor((1,), dtype="float32"),
+ in_2: R.Tensor((1, 1), dtype="float32"),
+ ) -> R.Tensor((1, 1), dtype="float32"):
+ R.func_attr({"num_input": 2})
+ with R.dataflow():
+ gv: R.Tensor((1, 1), dtype="float32") = R.broadcast_to(in_,
R.shape([1, 1]))
+ R.output(gv)
+ return gv
+
_assert_expand_ir("expand_with_dim_unchanged_test", [3, 1], [3, 4], [3,
4], ExpectedSameRank)
_assert_expand_ir("expand_with_diff_dim", [3, 1], [1, 3, 4], [1, 3, 4],
ExpectedHigherRank)
_assert_expand_ir(
@@ -5609,6 +5622,12 @@ def test_expand():
_assert_expand_dynamic_shapeexpr_ir(
"expand_with_dynamic_dim", [1, 32, 32], ["batch", 32, 32],
ExpectedDynamicShape
)
+ _assert_expand_dynamic_shapeexpr_ir(
+ "expand_with_higher_rank_same_padded_shape",
+ [1],
+ [1, 1],
+ ExpectedHigherRankSamePaddedShape,
+ )
def test_expand_incompatible_broadcasting():
