[Fix][Relax][ONNX] Preserve rank-expanding Expand

python/tvm/relax/frontend/onnx/onnx_frontend.py

@@ -2875,6 +2875,7 @@ def _impl_v13(cls, bb, inputs, attr, params):
         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 @@ def _impl_v13(cls, bb, inputs, attr, params):
                                 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))
 

tests/python/relax/test_frontend_onnx.py

@@ -5601,6 +5601,19 @@ def main(
                 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 main(
     _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():