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

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"):

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"