Add partial data propagation to enhance shape inference - #26269
Merged
Conversation
…ta propagation to work correctly
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
yuslepukhin
suggested changes
Oct 14, 2025
yuslepukhin
reviewed
Oct 14, 2025
yuslepukhin
reviewed
Oct 14, 2025
yuslepukhin
previously approved these changes
Oct 14, 2025
yuslepukhin
suggested changes
Nov 18, 2025
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
yuslepukhin
reviewed
Nov 19, 2025
yuslepukhin
reviewed
Nov 19, 2025
yuslepukhin
reviewed
Nov 19, 2025
yuslepukhin
reviewed
Nov 19, 2025
yuslepukhin
reviewed
Nov 19, 2025
yuslepukhin
reviewed
Dec 2, 2025
alex-spacemit
pushed a commit
to spacemit-com/onnxruntime
that referenced
this pull request
Dec 8, 2025
) ### Description Calling an operator's `TypeAndShapeInferenceFunction()` alone is sometimes insufficient for complete shape inference. For example, the `Shape` operator only infers the output’s rank (a 1-dimensional tensor) but not its actual dimension values. For instance, given an input of shape [1, 3, 64, 64], the Shape operator's `TypeAndShapeInferenceFunction()` produces an output shape tensor with 1-dimension as int64[4], representing the rank of the input tensor. Therefore, as you can imagine, the below graph's output shape can't be properly inferred (even though the input shape is known) because the shape data is lost at the `Shape `operator. <img width="563" height="488" alt="image" src="https://github.com/user-attachments/assets/bfa9fd8f-5291-4c6d-a679-3ce4a8c48669" /> To solve the issue, the `PartialDataPropagationFunction()`, defined in the ONNX operator schema, must also be executed to obtain the concrete output shape values, allowing accurate propagation of shape information throughout the graph. This PR adds the support of executing operator's `PartialDataPropagationFunction()` in ORT, and makes sure the shape values is properly propagated throughout the graph. ### Motivation and Context When using the Compile API to generate an EPContext model, all graph optimizations are disabled by default except for free dimension overrides. However, for certain models, such as a VAE decoder, the output shape may still fail to be properly inferred even when free dimension override values are provided beforehand. However, you won't hit this issue if enabling all the graph optimizations as some nodes, e.g. `Shape`, `Reshape `.. will be constant folded.
Sumit2318
pushed a commit
that referenced
this pull request
Jan 6, 2026
### Description Calling an operator's `TypeAndShapeInferenceFunction()` alone is sometimes insufficient for complete shape inference. For example, the `Shape` operator only infers the output’s rank (a 1-dimensional tensor) but not its actual dimension values. For instance, given an input of shape [1, 3, 64, 64], the Shape operator's `TypeAndShapeInferenceFunction()` produces an output shape tensor with 1-dimension as int64[4], representing the rank of the input tensor. Therefore, as you can imagine, the below graph's output shape can't be properly inferred (even though the input shape is known) because the shape data is lost at the `Shape `operator. <img width="563" height="488" alt="image" src="https://github.com/user-attachments/assets/bfa9fd8f-5291-4c6d-a679-3ce4a8c48669" /> To solve the issue, the `PartialDataPropagationFunction()`, defined in the ONNX operator schema, must also be executed to obtain the concrete output shape values, allowing accurate propagation of shape information throughout the graph. This PR adds the support of executing operator's `PartialDataPropagationFunction()` in ORT, and makes sure the shape values is properly propagated throughout the graph. ### Motivation and Context When using the Compile API to generate an EPContext model, all graph optimizations are disabled by default except for free dimension overrides. However, for certain models, such as a VAE decoder, the output shape may still fail to be properly inferred even when free dimension override values are provided beforehand. However, you won't hit this issue if enabling all the graph optimizations as some nodes, e.g. `Shape`, `Reshape `.. will be constant folded.
This was referenced Jun 16, 2026
titaiwangms
added a commit
that referenced
this pull request
Jun 23, 2026
…x output in data propagation (#29084) ## Summary A spec-valid `Shape → Gather(1-D index [-1]) → TopK` model fails to load since ORT 1.25.0 with: ``` K input must be a one-dimensional tensor of size 1. ``` The model is valid: a rank-1 (single-element) Gather index produces a rank-1 Gather output, so the value feeding TopK's `K` input is a 1-D size-1 tensor — exactly what TopK requires. The failure was an **ORT rank-preservation bug in shape-inference data propagation**, not a problem with the model. **Root cause.** `GatherOpDataPropagation::infer()` routed by element **count** rather than index **rank**: it guarded on `indices.size() == 1`, which is true for *both* a 0-D scalar index and a 1-D single-element index, and then unconditionally called `SetInferredShapeScalarValue()`. That dropped the rank of the spec-valid 1-D size-1 case, so `Graph::getInputData()` emitted a 0-D (dimensionless) propagated value. ONNX TopK shape inference then correctly rejected the 0-D `K`. This path was introduced by #26269 (partial data propagation to enhance shape inference). This reproduces even at `GraphOptimizationLevel.ORT_DISABLE_ALL`, where constant folding never runs — confirming the cause is data propagation in shape inference, **not** constant folding (#26345 was an earlier mis-attribution; see the corrected analysis). Fixes the regression reported in #29072. Corrected root-cause analysis: #29072 (comment) ## The fix - **Gather — rank-based routing.** Distinguish the index rank instead of its element count. A genuine 0-D scalar index still stores a scalar value; a rank-1 single-element index now stores a **rank-1** value, so `getInputData()` emits a `TensorProto` with `dims=[1]` and downstream TopK sees a valid 1-D size-1 `K`. The index rank is taken from the same constant initializer the index value comes from (via `get_initialized_input_values` now reporting the initializer rank), rather than a second, independently-resolved `NodeArg` shape — removing a potential source-of-truth drift (EDGE #2). - **Rank-tolerant elementwise companion (Add/Sub/Mul/Div).** These ops were scalar-only and would silently stop propagating once an operand became a rank-1 value (e.g. a `Shape → Gather(1-D idx) → Mul → TopK` chain), because the custom-propagation result replaces ONNX's rank-correct fallback. They now accept a single element carried as either a rank-0 scalar or a rank-1 `[1]` value and keep the output rank consistent with ONNX broadcasting (rank-1 if any operand is rank-1, else scalar), so such chains keep propagating end-to-end. Div additionally guards against division by zero. - **Shared helper (`data_propagation_value_utils.h`).** Centralizes reading/writing a single-element shape value while preserving its rank, used by both the Gather producer and the elementwise consumers so they cannot disagree on rank. The reader **declines** a rank-1 multi-element value (it must never collapse to `element[0]`), so a multi-element value can never be mistaken for a single one. ## Testing Five `ShapeInferenceV2Test` cases (with fixtures + generators), all loading the model at **every** optimization level (including `ORT_DISABLE_ALL`): - `GatherToTopKRankPreservationTest` — the core `Shape → Gather([-1]) → TopK` regression; asserts the rank-1 `K` is preserved. - `GatherMulToTopKRankPreservationTest` — the `… → Gather(1-D idx) → Mul → TopK` chain; asserts propagation survives the elementwise op. - `SinglePropagatedShapeValueGuardTest` — a direct unit test pinning the shared reader's behavior on each channel (scalar, rank-1 single-element, rank-1 multi-element, symbolic, empty). **Mutation-proven**: relaxing the `dim_size()==1` guard makes this test fail, restoring it makes it pass — so the guard the whole fix hinges on is test-locked. - `ShapeMulMultiElementNoScalarCollapseTest` — end-to-end check that a multi-element `Shape → Mul → ConstantOfShape` chain still resolves to its full rank-2 shape (no bogus scalar collapse). - `PartialDataPropagationTest` — pre-existing scalar-index coverage, unchanged. Full `onnxruntime_test_all` suite passes (0 failures) on top of the current `main` (opset-27 / ONNX 1.22.0 integration). The constant-folding memory path (#26345) is untouched — the diff is confined to `data_propagation/`, a small `graph.cc` change, and tests. ## Follow-ups (intentionally out of scope for this PR) - Hardening for a rank ≥ 2 single-element index (e.g. shape `[1,1]`) to *decline* rather than route as rank-1 — needs its own discriminating unit test; pathological/non-exporter, worst case is degraded inference rather than a crash. - Explicit end-to-end coverage for Add/Sub/Div rank-1 chains (the shared-reader unit test already covers the read path for all four ops; only Mul is currently exercised end-to-end). - Minor readability nits. ## DCO Commit is DCO signed-off. --------- Signed-off-by: titaiwangms <titaiwang@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Calling an operator's
TypeAndShapeInferenceFunction()alone is sometimes insufficient for complete shape inference.For example, the
Shapeoperator only infers the output’s rank (a 1-dimensional tensor) but not its actual dimension values.For instance, given an input of shape [1, 3, 64, 64], the Shape operator's
TypeAndShapeInferenceFunction()produces an output shape tensor with 1-dimension as int64[4], representing the rank of the input tensor.Therefore, as you can imagine, the below graph's output shape can't be properly inferred (even though the input shape is known) because the shape data is lost at the
Shapeoperator.To solve the issue, the
PartialDataPropagationFunction(), defined in the ONNX operator schema, must also be executed to obtain the concrete output shape values, allowing accurate propagation of shape information throughout the graph.This PR adds the support of executing operator's
PartialDataPropagationFunction()in ORT, and makes sure the shape values is properly propagated throughout the graph.Motivation and Context
When using the Compile API to generate an EPContext model, all graph optimizations are disabled by default except for free dimension overrides. However, for certain models, such as a VAE decoder, the output shape may still fail to be properly inferred even when free dimension override values are provided beforehand.
However, you won't hit this issue if enabling all the graph optimizations as some nodes, e.g.
Shape,Reshape.. will be constant folded.