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 d02a68e403 [Relax][Frontend][ONNX] Support dynamic index for Gather on
shape (#19968)
d02a68e403 is described below
commit d02a68e403eb1b473844bd1b36302f44ac66dc32
Author: Hamza Qureshi <[email protected]>
AuthorDate: Sat Jul 18 10:10:32 2026 +0500
[Relax][Frontend][ONNX] Support dynamic index for Gather on shape (#19968)
The ONNX importer's Gather converter asserted that indices must be a
constant whenever the data operand is a ShapeExpr, raising "Only
constant indices supported for shape gather." for any runtime-computed
index. Detection post-processing graphs such as FasterRCNN feed a
dynamic index into a Gather whose data comes from a Shape node, so
import failed before compilation could start.
Keep the fast path for a single constant index, which resolves one
dimension to a PrimValue and preserves shape-specialized handling
downstream. Any other index (dynamic, or a constant selecting multiple
dimensions) materializes the shape as an int64 tensor via
shape_to_tensor and gathers from it at runtime, reusing the existing
negative-index normalization.
Adds a regression test that gathers a dimension out of a Shape result
using a non-constant index, covering positive and negative indices, and
checks it against onnxruntime.
Fixes part of #19965.
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 25 ++++---
tests/python/relax/test_frontend_onnx.py | 93 +++++++++++++++++++++++++
2 files changed, 108 insertions(+), 10 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 9121018bb9..881789c7ed 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -1333,17 +1333,22 @@ class Gather(OnnxOpConverter):
output = _np.take(data.data.numpy(), indices.data.numpy(),
axis=axis)
return relax.const(output, output.dtype)
- # If input is a shape expression, take a value from that shape and
return it as a constant.
+ # If input is a shape expression, take a value from that shape. A 0-D
+ # scalar constant index resolves to one dimension that we return as a
+ # PrimValue to keep shape-specialized handling in downstream
+ # shape-construction patterns. Any other index materializes the shape
as
+ # an int64 tensor and gathers from it at runtime, reusing the
+ # negative-index normalization below. ONNX Gather defines the output
rank
+ # as q + r - 1 (q = rank of indices); since the shape is rank 1, a
+ # non-scalar index such as (1,) must keep its rank, so only a true
+ # 0-D index collapses to a scalar.
if isinstance(data, relax.ShapeExpr):
- assert isinstance(indices, relax.Constant), (
- "Only constant indices supported for shape gather."
- )
- np_index = indices.data.numpy()
- if len(np_index.shape) == 1:
- np_index = np_index[0]
- np_index = int(np_index)
- shape_val = data[np_index]
- return relax.prim_value(shape_val)
+ if isinstance(indices, relax.Constant) and
indices.data.numpy().ndim == 0:
+ np_index = int(indices.data.numpy().item())
+ shape_val = data[np_index]
+ return relax.prim_value(shape_val)
+
+ data = bb.normalize(relax.op.shape_to_tensor(data))
indices_dtype = indices.ty.dtype.dtype
if not indices_dtype.startswith("uint"):
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 4a5c2f778d..6b449624b2 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -1430,6 +1430,99 @@ def test_gather():
_verify_gather([3, 3], [[0, 2]], [3, 1, 2], ExpectedRank2Axis1, 1)
[email protected]("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.
+
+ Detection post-processing graphs (e.g. FasterRCNN) feed a runtime-computed
+ index into a Gather whose data is a Shape output. The index is not a
+ constant, so the importer must materialize the shape as a tensor and gather
+ from it at runtime rather than resolving the dimension at compile time.
+ """
+ data_shape = [3, 4, 5, 6]
+ shape_node = helper.make_node("Shape", ["data"], ["shape"])
+ gather_node = helper.make_node("Gather", ["shape", "index"], ["y"], axis=0)
+
+ graph = helper.make_graph(
+ [shape_node, gather_node],
+ "gather_shape_dynamic_index_test",
+ inputs=[
+ helper.make_tensor_value_info("data", TensorProto.FLOAT,
data_shape),
+ helper.make_tensor_value_info("index", TensorProto.INT64, []),
+ ],
+ outputs=[helper.make_tensor_value_info("y", TensorProto.INT64, [])],
+ )
+
+ model = helper.make_model(graph,
producer_name="gather_shape_dynamic_index_test")
+ input_values = {
+ "data": np.random.randn(*data_shape).astype("float32"),
+ "index": np.array(index).astype("int64"),
+ }
+ check_correctness(model, inputs=input_values)
+
+
[email protected](
+ "indices",
+ [
+ np.array(2, dtype="int64"), # 0-D scalar -> PrimValue fast path
+ np.array([2], dtype="int64"), # (1,) -> must stay rank 1
+ np.array([[2]], dtype="int64"), # (1, 1) -> must stay rank 2
+ np.array([1, 3], dtype="int64"), # (2,) -> multiple dims
+ np.array([-1], dtype="int64"), # (1,) negative index
+ ],
+)
+def test_gather_shape_constant_index(indices):
+ """Gather from a Shape result using a constant index of varying rank.
+
+ ONNX Gather defines the output rank as q + r - 1 where q is the rank of the
+ indices. Since a Shape output is rank 1, only a true 0-D scalar index
should
+ collapse to a scalar; a (1,) index must produce a rank-1 result rather than
+ being folded into a PrimValue.
+ """
+ data_shape = [3, 4, 5, 6]
+ shape_node = helper.make_node("Shape", ["data"], ["shape"])
+ # Emit the indices through a Constant node: an initializer would become a
+ # function parameter under keep_params_in_input=True and bypass the
+ # constant fast path in the Gather converter.
+ const_node = helper.make_node(
+ "Constant",
+ [],
+ ["indices"],
+ value=helper.make_tensor(
+ "value",
+ TensorProto.INT64,
+ indices.shape,
+ indices.flatten().tolist(),
+ ),
+ )
+ gather_node = helper.make_node("Gather", ["shape", "indices"], ["y"],
axis=0)
+
+ graph = helper.make_graph(
+ [shape_node, const_node, gather_node],
+ "gather_shape_constant_index_test",
+ inputs=[
+ helper.make_tensor_value_info("data", TensorProto.FLOAT,
data_shape),
+ ],
+ outputs=[helper.make_tensor_value_info("y", TensorProto.INT64,
list(indices.shape))],
+ )
+
+ model = helper.make_model(graph,
producer_name="gather_shape_constant_index_test")
+ input_values = {
+ "data": np.random.randn(*data_shape).astype("float32"),
+ }
+ check_correctness(model, inputs=input_values)
+
+ # check_correctness broadcasts a scalar against a one-element tensor, so
+ # assert the rank explicitly: only a 0-D index may collapse to a scalar.
+ tvm_out = run_in_tvm(model, inputs=input_values)
+ if isinstance(tvm_out, tvm.runtime.Tensor):
+ out_shape = tuple(tvm_out.numpy().shape)
+ else:
+ # PrimValue fast-path outputs come back as plain Python scalars.
+ out_shape = ()
+ assert out_shape == indices.shape
+
+
def _make_gather_negative_indices_expected(axis: int, indices_shape,
indices_type):
indices_shape = tuple(indices_shape)
indices_dtype = "int64" if indices_type == TensorProto.INT64 else "int32"
