Skip to content

Dynamic shape explainer - #225

Draft
theHamsta wants to merge 13 commits into
rustnn:mainfrom
theHamsta:dynamic-shape-explainer
Draft

Dynamic shape explainer#225
theHamsta wants to merge 13 commits into
rustnn:mainfrom
theHamsta:dynamic-shape-explainer

Conversation

@theHamsta

Copy link
Copy Markdown
Collaborator

Summary

  • Describe the user-visible and code-level changes.

WIP implementation of webmachinelearning/webnn#945 basically adding the required API functions.

The .shape/.shape_with_options function existed already before this PR and where not guarded by feature = "dynamic-inputs. The shape op currently returns Vec<i64> instead of the u32 from the explainer. ORT and we used i64 internally but should probably switch to u32 from WebNN.

This PR will require some kind of shape validation. Either

  • our existing, which may give up on dynamic shape resolution (e.g. yield unknown shape),
  • or no build validation for dynamic shapes
  • or a simple expression tree that is evaluated on computeShape
  • or a fully symbolic constraint solver like z3 which is likely overkill but might help us for bring up by being mathematically correct and ready to use

Validation

  • make test
  • Relevant WPT or integration checks

Documentation

  • Updated docs if behavior changed
  • If backend converter/executor operator support changed, ran make docs-backend-ops and committed docs/development/backend-operator-support.md

@theHamsta
theHamsta force-pushed the dynamic-shape-explainer branch from 5386fb2 to b8028a1 Compare September 10, 2026 10:36
@matthargett

Copy link
Copy Markdown
Contributor

I've updated #223 against main and added exact-result tests for shape -> unsqueeze and integer-mask where. A few integration points from the CoreML work:

  • I'd keep the existing int64 interchange path working while introducing the explainer's uint32 API contract. MIL shape returns int32, so this needs an explicit, checked backend boundary, not a blanket dtype change or feature gate on imported shape nodes.
  • The latest clarification makes a small, shared Rust shape evaluator look like the useful next step: input shapes and selected constants only, with the same validation before dispatch. It is a side computation; the backend still receives the shape operators. It won't by itself fix the CoreML lowering gaps in Track SmolLM-135M CoreML blockers #222.
  • A reduced shape -> unsqueeze graph with input [1, sequence<=8, 4] still hits the model-description/MIL-input shape mismatch when compiling on my M4. The Watch exporter work also showed why unknown dimensions, concrete defaults/ranges, and logical zero extents must stay distinct. In our CPU-only SE2/watchOS 11.6.2 runs, empty concat works with non-null backing; that isn't permission to replace zero with one or skip arbitrary empty operations (device results). Shape-product/byte-size arithmetic also needs checked conversions on arm64_32.
  • MLResample2dDynamicOptions currently lacks the sizes operand. Could we add that, track it as an input dependency, and preserve axes/scales/sizes through interchange? That keeps Complete CoreML resample2d lowering #215's arbitrary-axis behavior intact when exchanging WebNN graphs.

Would tests comparing the host shape result with actual CoreML outputs be useful as the next shared validation step? I can help with those.

mtavenrath pushed a commit that referenced this pull request Sep 13, 2026
Gather index normalization currently replaces dynamic dimensions with
their maximum extents. The model interface and dispatch also use those
maxima, preventing active-sized inputs from reaching CoreML correctly.

This preserves dynamic dimensions through `gather`, `gatherElements`,
and `gatherND` normalization, adds matching model shape ranges
(including integer proxy outputs), and binds the active input dimensions
with checked byte lengths. Dispatch now rejects oversized results
instead of silently truncating them.

Scalar gathers also bypassed normalization: gathering index `-1` from
`[10, 20, 30]` returned `10` instead of `30`. Normalize constant and
runtime scalar indices while preserving WebNN rank-zero results and
CoreML's `[1]` feature boundary, consistent with
[webmachinelearning/webnn#953](webmachinelearning/webnn#953).

Related to [#222](#222). Dynamic
behavior remains behind `dynamic-inputs`; this complements the proposed
host-side shape evaluation in
[#225](#225). Dynamic gathered-axis
bounds, runtime reshape/slice lowering, and empty KV-cache behavior
remain separate work.

Validation on the M4:

- Formatting and strict Clippy passed; 356 default, 364 CoreML, and 376
CoreML/dynamic Rust tests passed.
- Locally compiled CoreML models returned exact values for all three
gather variants over active lengths `1 -> 4 -> 2 -> 1`, including
negative/out-of-range indices; constant and runtime scalar cases passed.
Converter tests cover mixed dimensions, axis removal, and integer proxy
metadata.
- CoreML WPT: 2,461 passed, 40 skipped, and the same 20 negative-index
scatter failures as clean `main` (`04fc5f4`), using [WPT
`04cd4a8`](web-platform-tests/wpt@04cd4a8).
The previously skipped scalar out-of-range positive gather case now
passes; no expectations or snapshot changes.

No newer CoreML operations or APIs are introduced. Physical iOS
18.x/watchOS 11.x validation has not been rerun for this patch.
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