From 44b2dff833f7a6893ee2e0e41bc3accec5a96785 Mon Sep 17 00:00:00 2001 From: Hamza Qureshi Date: Wed, 8 Jul 2026 17:42:04 +0500 Subject: [PATCH] [Relax][Frontend][ONNX] Support dynamic index for Gather on shape 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. --- .../tvm/relax/frontend/onnx/onnx_frontend.py | 25 +++++++++------ tests/python/relax/test_frontend_onnx.py | 31 +++++++++++++++++++ 2 files changed, 46 insertions(+), 10 deletions(-) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index f46041831525..2145b7e26390 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -1182,17 +1182,22 @@ def _impl_v13(cls, bb, inputs, attr, params): 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 single + # 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 (dynamic, or a constant selecting multiple dimensions) + # materializes the shape as an int64 tensor and gathers from it at runtime, + # reusing the negative-index normalization below. 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().size == 1: + 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) + + 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 9058cdf70ae7..c3830cf7121c 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -1462,6 +1462,37 @@ def main( _verify_gather([3, 3], [[0, 2]], [3, 1, 2], ExpectedRank2Axis1, 1) +@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. + + 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) + + 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"