Shape of Gather output is wrong making it unusable as K input to TopK operator

[arnej27959]

Describe the issue

After upgrading from onnxruntime 1.23 to 1.26 one of our customers models fails to load with this error:

Load model from topk_repro.onnx failed:Node (/TopK) Op (TopK) [ShapeInferenceError] K input must be a one-dimensional tensor of size 1.

As far as I can tell this happens because there's a "Gather" operation that should give the required shape as input, but doesn't. I've reduced the model to just a toy example; the "Gather" operation here should be no-op but clearly its output type isn't as expected.

topk.py

To reproduce

on both Mac and Linux I can reproduce the problem with the attached script:

python3.14 -m venv ~/onnxruntime-env
source ~/onnxruntime-env/bin/activate
pip install onnx onnxruntime==1.26.0
python topk.py

Urgency

No response

Platform

Linux

OS Version

AlmaLinux 10.2

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

v1.26.0

ONNX Runtime API

Python

Architecture

X64

Execution Provider

Default CPU

Execution Provider Library Version

No response

[damdoo01-arm]

Attempting fix in ONNX itself:
onnx/onnx#8090

[titaiwangms]

⚠️ CORRECTION (edited): the root-cause attribution below is incorrect. Further investigation — reproducing the model at GraphOptimizationLevel.ORT_DISABLE_ALL, where constant folding does not run, and the failure still occurs — showed this is not caused by PR #26345 (constant folding). The actual regressing change is PR #26269 ("Add partial data propagation to enhance shape inference"), specifically GatherOpDataPropagation::infer() scalarizing a 1-D single-element Gather output. See the full corrected analysis here: #29072 (comment) — the original text is preserved below for the record.

---

Summary

This TopK load-time [ShapeInferenceError] K input must be a one-dimensional tensor of size 1 is an ONNX Runtime regression introduced during constant folding, not an ONNX bug. The ONNX TopK shape-inference check is unchanged since opset 10 (2019) and is spec-compliant. The model in the repro is itself spec-valid; the failure comes from ORT materializing the folded K as a 0-D scalar.

Regressing change

The behavior was introduced by PR #26345 ("Save much memory at model loading time by converting weights to OrtValues early", commit b86098800f), first shipped in ONNX Runtime 1.25.0.

Version-window evidence (release membership of b86098800f):

• v1.23.0 (2025-09-26) — does not contain it → model loads (matches the report).
• v1.25.0 (2026-04-20) — contains it.
• v1.26.0 (2026-05-08) — contains it → model fails to load (matches the report).
• Note there is no v1.24.0 release; tags go 1.23 → 1.25 → 1.26.

Other initializer/OrtValue changes were ruled out: PRs #23979 and #25652 merged before the 1.23 branch cut (so they shipped in 1.23, which works), and #27620 is a later crash fix that is not an ancestor of 1.25.0.

Mechanism

In the repro the chain is Shape(X) -> Gather(S, indices=[-1]) -> TopK(X, K):

• The Gather index is the 1-D constant [-1], so the spec-correct Gather output is a 1-D, size-1 tensor ([2000], shape [1]), which satisfies TopK's requirement that K be a 1-D tensor of size 1.
• PR Save much memory at model loading time by converting weights to OrtValues early #26345 reroutes constant-folded outputs through an OrtValue ⇄ TensorProto round-trip (constant_folding.cc now uses TensorToTensorProto(..., use_tensor_buffer=true) and AddInitializedOrtValue; graph.cc adds the early put_large_tensor_in_ort_value conversion).
• For this single-element folded tensor, that round-trip drops the rank-1 [1] shape and materializes K as a 0-D scalar initializer.
• ONNX TopK shape inference then correctly rejects the 0-D K, because the spec defines K as "a 1-D tensor containing a single positive value..."

In short: a spec-valid 1-D size-1 K is being reduced to a 0-D scalar during folding, which is what trips the (correct) ONNX check.

Confidence

High, based on release-membership and code-path analysis. The exact line that emits dims [] instead of [1] for a single-element folded tensor would need an ORT build plus git bisect v1.23.0..v1.25.0 against the repro to confirm to the commit; #26345 is the predicted first-bad commit (the value/shape is computed correctly in ConstantFolding::ApplyImpl, so the drop most likely occurs inside the TensorToTensorProto(..., use_tensor_buffer=true) / OrtValue round-trip).

Suggested fix direction

ORT constant folding should preserve the rank-1 shape of single-element folded tensors rather than collapsing them to a 0-D scalar, so that a folded Shape -> Gather K remains a valid 1-D size-1 input to TopK.

Note on the ONNX side (PR #8090)

There is a proposed ONNX change, onnx/onnx#8090, that would relax the TopK K shape check (from "rank-1 {1}" to "exactly one element", which would also accept a 0-D scalar). That is a reasonable defensive, ecosystem-compatibility improvement, but it is not the root fix for this regression, and it is currently incomplete: its data-propagation test path is failing in CI, and a complete change would also need the bundled reference implementation (onnx/reference/ops/op_topk.py) normalized to accept a single-element K (it currently rejects a 0-D K as well). The root cause remains the ORT-side rank drop described above.

[hariharans29]

@titaiwangms - Thanks for this finding ! It slipped my mind that (as rare as it is) something from the optimizer suite could have created an intermediate invalid model.

[titaiwangms]

Correction / erratum

My earlier comment above mis-attributed the root cause to constant folding (ORT PR #26345). That attribution was based on release-window plus code-path analysis without a build, and it is incorrect. Correcting the record below; the rest of the earlier comment's conclusions (the model is spec-valid, the ONNX TopK check is spec-correct, and ONNX PR #8090 is a defensive-but-incomplete compatibility net) still stand.

Decisive evidence that exonerates #26345

Reproducing the reported model (Gather index is the 1-D constant [-1]) against released wheels while toggling graph optimization:

• On 1.26.0, the model fails to load even at GraphOptimizationLevel.ORT_DISABLE_ALL.
• Constant folding (Save much memory at model loading time by converting weights to OrtValues early #26345) only runs at ORT_ENABLE_BASIC and above, so it cannot be responsible for a failure that reproduces with all graph optimizations disabled.

The failure therefore happens earlier, in shape-inference data propagation during Graph::Resolve, not in constant folding.

Actual regressing change

PR #26269 — "Add partial data propagation to enhance shape inference" (merge commit a2c4374fe7b6, merged 2025-12-04). It introduced the onnxruntime/core/graph/data_propagation/ framework, including gather_op_data_propagation.cc.

Same version window as before:

• v1.23.0 — does not contain Add partial data propagation to enhance shape inference #26269 → model loads.
• v1.25.0 / v1.26.0 — contain Add partial data propagation to enhance shape inference #26269 → model fails to load.

This also matches the description of "computing the K value before runtime": that is the new data-propagation feature, not the memory optimization in #26345.

Root cause

In GatherOpDataPropagation::infer() (gather_op_data_propagation.cc):

• The guard is indices.size() == 1, i.e. the element count of the indices tensor. That is 1 for both a 0-D scalar index (-1, output correctly rank-0) and a 1-D single-element index ([-1], output should remain rank-1 [1]).
• In both cases it calls SetInferredShapeScalarValue(...), which stores the result as a scalar and drops the rank for the 1-D case.
• Graph::getInputData(...) then materializes that scalar as a TensorProto carrying int64_data but no dims (a 0-D tensor) — matching the "TensorProto with no dimensions" described earlier in this thread.

ONNX TopK shape inference then correctly rejects the 0-D K, because the spec requires K to be a 1-D tensor of size 1. So a spec-valid 1-D size-1 K is being scalarized during data propagation, which is what trips the (correct) ONNX check.

Suggested fix direction

In GatherOpDataPropagation::infer(), distinguish the indices rank rather than the element count: keep the scalar store only for a genuine 0-D scalar index, and route the rank-1 single-element case to the rank-preserving path (inferred_shape_values_, which records dims). This is general — it keeps the rank-1 shape of any Shape -> Gather single-element result for all downstream consumers (TopK, Reshape, Slice, Expand, …), not just TopK.

On the ONNX side (PR #8090)

onnx/onnx#8090 proposes relaxing the ONNX TopK K check (rank-1 {1} → exactly-one-element, which would also accept a 0-D scalar). That remains a reasonable defensive, ecosystem-compatibility improvement, but it is not the root fix for this regression, and it is currently incomplete (its data-propagation test path is failing in CI, and a complete change would also need the bundled reference implementation onnx/reference/ops/op_topk.py normalized to accept a single-element K). The root cause is the ORT-side rank drop described above.

[chilo-ms]

Thanks @titaiwangms for the PR and i left some comments. It should fix the issue of failing to run shape inference on TopK op in the repro model.

[arnej27959]

I can confirm that recompiling 1.26.0 adding #29084 as a patch fixes this problem for me.