ONNX ReferenceEvaluator rejects valid GatherElements subshapes

Finding

GatherElements requires data and indices to have equal rank, and defines the output shape as the shape of indices. Coordinates on each non-gather axis range over the corresponding index domain; those dimensions may therefore be smaller than the data dimensions.

The released ReferenceEvaluator imposes a stricter condition: every non-gather dimension must equal the corresponding data dimension. It rejects valid subshape inputs before performing any gather.

Minimal counterexample

data = [[10, 11, 12],
        [20, 21, 22]]              shape (2, 3)
indices = [[2, 0]]                 shape (1, 2)
axis = 1

ReferenceEvaluator: ValueError requiring equal non-axis dimensions
direct indexing:    [[12, 10]]
ONNX Runtime:       [[12, 10]]

The first output row refers only to the first data row, so the smaller leading dimension is well-defined and accepted by ONNX Runtime.

Verification

• Reproduced with ONNX 1.22.0, ONNX Runtime 1.29.0, NumPy 2.5.2, Python 3.12 and the CPU execution provider.
• Confirmed unchanged on ONNX main commit c71cbb485c97aa2cb257626a5ac5b0f205c49a22.
• Three regression cases cover axes 0, 1 and -1 in two and three dimensions.
• Thirty-six differential cases cover equal and smaller non-axis dimensions, positive and negative axes, and positive and negative indices; all corrected outputs match ONNX Runtime exactly.
• The complete reference-evaluator test module passed: 454 tests, 4 skipped.
• Restoring the released implementation makes all three new tests fail with the original shape-equality exception.
• Ruff format/check and git diff --check pass.
• Searches of ONNX issues and pull requests found no matching report.

Correction

Restrict data to the coordinate extent of the index domain on non-gather axes, then apply numpy.take_along_axis. This directly models the operator equations, preserves equal-shape behavior, and supports positive and negative indices.

The correction and regression tests are submitted in ONNX PR #8444. The public reproducer is in gero-onnx-gatherelements-subshape-audit.

Boundary

This is a reference-implementation correctness defect, not a security issue. It affects validation of otherwise valid GatherElements models. ONNX Runtime produces the specified output for the counterexample.