Skip to content

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

Open
hamzaqureshi5 wants to merge 1 commit into
apache:mainfrom
hamzaqureshi5:relax-onnx-gather-shape-dynamic-index
Open

[Relax][Frontend][ONNX] Support dynamic index for Gather on shape#19968
hamzaqureshi5 wants to merge 1 commit into
apache:mainfrom
hamzaqureshi5:relax-onnx-gather-shape-dynamic-index

Conversation

@hamzaqureshi5

Copy link
Copy Markdown

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.

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.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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

@tlopex tlopex left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants