From 44b2dff833f7a6893ee2e0e41bc3accec5a96785 Mon Sep 17 00:00:00 2001 From: Hamza Qureshi Date: Wed, 8 Jul 2026 17:42:04 +0500 Subject: [PATCH 1/2] [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" From 21f253b68755e6b85a29dd4a800909f72779ad7c Mon Sep 17 00:00:00 2001 From: Hamza Qureshi Date: Fri, 17 Jul 2026 15:43:33 +0500 Subject: [PATCH 2/2] [Relax][Frontend][ONNX] Restrict shape-gather constant fast path to 0-D indices Address review: ONNX Gather defines the output rank as q + r - 1 where q is the rank of indices, so with rank-1 shape data only a true 0-D scalar constant index may collapse to a PrimValue. The previous size == 1 check wrongly folded (1,)- and (1, 1)-shaped constant indices to a scalar; they now go through the shape_to_tensor path and keep their rank. Adds a parametrized test covering constant indices of rank 0, 1, and 2. The indices are emitted via a Constant node (an initializer would become a function parameter under keep_params_in_input=True and bypass the fast path), and the output rank is asserted explicitly since assert_allclose broadcasts a scalar against a one-element tensor. --- .../tvm/relax/frontend/onnx/onnx_frontend.py | 22 +++---- tests/python/relax/test_frontend_onnx.py | 62 +++++++++++++++++++ 2 files changed, 73 insertions(+), 11 deletions(-) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index f7917bfbf6e2..88d06b216895 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -1333,18 +1333,18 @@ 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. 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 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): - 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) + 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) diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index b6ae9c258f59..cd8cb2285f52 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -1461,6 +1461,68 @@ def test_gather_shape_dynamic_index(index): check_correctness(model, inputs=input_values) +@pytest.mark.parametrize( + "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"