[Relax][Frontend][ONNX] Support dynamic index for Gather on shape#19968
[Relax][Frontend][ONNX] Support dynamic index for Gather on shape#19968hamzaqureshi5 wants to merge 1 commit into
Conversation
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 apache#19965.
There was a problem hiding this comment.
Code Review
This pull request updates the ONNX frontend's Gather operator implementation to support dynamic or multi-dimensional indices when gathering from a shape expression. Instead of asserting that indices must be constant, it now materializes the shape as an int64 tensor at runtime when necessary. A test case has been added to verify this behavior. The review feedback suggests optimizing the constant index extraction path by avoiding redundant .numpy() calls, preventing potential TypeError on higher-dimensional single-element constants, and simplifying the scalar extraction using .item().
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
tlopex
left a comment
There was a problem hiding this comment.
The constant fast path should be selected based on the rank of indices, rather than its number of elements.
ONNX Gather defines the output rank as q + r - 1, where q is the rank of indices. Since the Shape output is rank 1 here, the Gather output must preserve the rank of indices. A constant index with shape (1,) should therefore produce a tensor with shape (1,), not a scalar PrimValue.
I reproduced the following behavior with this PR:
indices shape [] -> scalar 4
indices shape [1] -> scalar 4 # expected [4]
indices shape [1, 1] -> scalar 4 # expected [[4]]
indices shape [2] -> [3, 5]
Could we restrict the PrimValue fast path to true 0-D scalar constants?
if isinstance(indices, relax.Constant):
np_indices = indices.data.numpy()
if np_indices.ndim == 0:
np_index = int(np_indices.item())
return relax.prim_value(data[np_index])
All non-scalar indices should use the shape_to_tensor path, even when they contain only one element. The suggested ndim <= 1 condition in the existing bot comment would still collapse a (1,) tensor incorrectly.
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.