diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index eaad12752499..88d06b216895 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/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"): diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index db9bf18a8110..cd8cb2285f52 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/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"