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 c985542951 [Fix][Relax][Frontend][ONNX] Support Shape outputs as
Gather indices (#20179)
c985542951 is described below
commit c985542951a1e85a6ced3b93acf6f6f11ff31148
Author: Gunse11er <[email protected]>
AuthorDate: Fri Aug 28 15:26:43 2026 +0800
[Fix][Relax][Frontend][ONNX] Support Shape outputs as Gather indices
(#20179)
The Relax ONNX frontend preserves Shape outputs as ShapeExpr values.
Gather assumed its indices operand always had a tensor type and accessed
its dtype, which raised an AttributeError when a Shape output was used
as indices.
Materialize ShapeExpr indices as int64 tensors before the existing dtype
and negative-index handling. Add an execution test that compares the
imported model against ONNX Runtime.
Fixes #20176.
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 3 +++
tests/python/relax/test_frontend_onnx.py | 27 +++++++++++++++++++++++++
2 files changed, 30 insertions(+)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index f1cf2957c6..024b2bf2b6 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -1361,6 +1361,9 @@ class Gather(OnnxOpConverter):
data = bb.normalize(relax.op.shape_to_tensor(data))
+ if isinstance(indices, relax.ShapeExpr):
+ indices = bb.normalize(relax.op.shape_to_tensor(indices))
+
indices_dtype = indices.ty.dtype.dtype
if not indices_dtype.startswith("uint"):
data_shape = bb.normalize(relax.op.shape_of(data))
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 9d1a6d8aaa..91b2e2ded7 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -1566,6 +1566,33 @@ def test_gather():
_verify_gather([3, 3], [[0, 2]], [3, 1, 2], ExpectedRank2Axis1, 1)
+def test_gather_indices_from_shape():
+ """Gather from a tensor using the dimensions of another tensor as
indices."""
+ shape_node = helper.make_node("Shape", ["shape_source"], ["indices"])
+ gather_node = helper.make_node("Gather", ["data", "indices"], ["y"],
axis=0)
+
+ graph = helper.make_graph(
+ [shape_node, gather_node],
+ "gather_indices_from_shape_test",
+ inputs=[
+ helper.make_tensor_value_info("data", TensorProto.FLOAT, [4]),
+ helper.make_tensor_value_info("shape_source", TensorProto.FLOAT,
[2, 3]),
+ ],
+ outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [2])],
+ )
+
+ model = helper.make_model(
+ graph,
+ producer_name="gather_indices_from_shape_test",
+ opset_imports=[helper.make_opsetid("", 18)],
+ )
+ input_values = {
+ "data": np.random.randn(4).astype("float32"),
+ "shape_source": np.random.randn(2, 3).astype("float32"),
+ }
+ check_correctness(model, inputs=input_values, opset=18)
+
+
@pytest.mark.parametrize("index", [0, 2, 3, -1, -4])
def test_gather_shape_dynamic_index(index):
"""Gather a dimension out of a Shape result using a non-constant index.
