[Relax][Frontend][ONNX] Support dynamic index for Gather on shape

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

@@ -1333,17 +1333,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 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"):

tests/python/relax/test_frontend_onnx.py

@@ -1430,6 +1430,99 @@ 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)
+
+
+@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"