Add dynamic shape explainer

[miaobin]

This is the initial draft to summarize the discussion on dynamic shapes #883 .

The explainer covers named/unnamed dynamic dimensions, deferred (dispatch-time) shape validation, computeShapes() API, and a new family of shape-as-data (*Dynamic) operators.

Open questions and considered alternatives are called out explicitly in the explainer and feedback on this doc would be very welcome.

[miaobin]

A POC based on the current explainer is WIP. And it has been validated against some real-world Transformer, LLM and image generation models.
With credits to @Honry for the ORT WebNN EP POC and assistance with the explainer, @huningxin and @fdwr for the initial review of this explainer.

[fdwr]

👍 I have some thoughts, but it's 95% 👌.

[fdwr]

Do you think we even need unnamed dimensions? A uniquely named dynamic dimension and an unnamed dynamic dimension are identical after all, and it's preferable to have debuggable symbols. The only reason to support them would be because existing callers may have them (like ORTWeb calling WebNN), but we could always just synthesize a name on the fly like "inputTensorName2_axis3" 🤔. I mean, if I was debugging and hit a shape inference error, I'd rather see that than just null for a name. Alternately using "" instead of null could be less problematic (no need to check for null first before trying to use/print the string).

Actually, seeing generated names would probably help you too during WebNN/Chromium debugging, seeing where pass-through fails during shape inference.

[miaobin]

Some history, since we did have derived names early on and removed them deliberately. We generated things like "height-3" and "broadcast_1_hight". The names weren't unique and two unrelated conv/slice calls would collide on the same name and the backend would conflate two independent dimensions. We patched that with a per-builder counter. Additionally, when we created the POC on the ORT backend at the begining, we only called OrtApi::SetDimensions. Regardless of the name, any dim that was dynamic was ultimately set to -1. Then we dropped derived names entirely in the rewrite and narrowed the rule to "a name survives only on 1:1 pass-through".

A synthesized name that is unique and implies no relationship is semantically identical to an anonymous dim, and it's strictly better to look at while debugging. I agree with the direction and will take it.

[miaobin]

Done, the explainer now requires every dynamic dimension to carry a name, and the unnamed state is gone. MLInputOperandDescriptor.shape drops its nullable elements entirely (sequence<MLDimension>), so null is no longer accepted at all and the empty string is rejected too.

@huningxin Do you have any thoughts on this point?

[w3cbot]

anssiko marked as non substantive for IPR from ash-nazg.

[anssiko]

I cleared the automatic IPR check for this PR, since explainer documents are consider non-substantive from the W3C Patent Policy point of view.

Details

This IPR check is to ensure normative portions of the specification come from organizations who participate the WG. Furthermore, in this case, I can attest the authors of this PR are affiliated with Intel, and as such, any normative portions are reusable in the spec PR, as appropriate.

[anssiko]

@reillyeon you had good questions comments on our last call for this. Do you have some other Googlers in mind who should review this explainer PR?

When this PR lands, the team will start landing the implementation in smaller chunks for further validation of this approach.

[reillyeon]

Overall I think this approach is reasonable. My main concern is that by making shapes dynamic there are more opportunities for memory safety issues in the implementation.

[anssiko]

With two approvals, editors are welcome to merge this explainer PR at will.

Thank you @miaobin for this contribution and Reilly, Dwayne, everyone for your review.

[huningxin]

Re "(and skipping oversized constants)", what are the oversized constants? Do you mean its length is over maximum rank limit? Should we reject rather than skip it?

[miaobin]

This is actually an implementation level design choice. To prevent the collection of large constants, we limit the maximum number of elements. Here, "skip" refers to the practice of not collecting a constant when an oversized one is encountered during the collection phase.

The description of this passage has been refined.

[huningxin]

we limit the maximum number of elements

What's the limit? Do you mean validate against valid dimension?

[miaobin]

What's the limit? Do you mean validate against valid dimension?

The limit is the rank of a constant. A constant with more elements than the limit is not a shape related vector, it is a weight, and we will not copy it.

[fdwr]

👍 Cool.

[miaobin]

@huningxin The latest Patchset of the dynamic shape POC runs dynamic shape inference and validation in Blink at dispatch time and also makes MLGraph.computeShapes() synchronous. PTAL, thanks.

[theHamsta]

the post-compilation on dispatch might be finger printable (below discussion suggest computeShapes is the natural place for post-compilation and shape specialization). Also, the implementation hitting an very unfortunate shape might have a finger printable cost that allows guesses about the used HW and implementation.

[theHamsta]

computeShapes() gives an implementation a natural place to do that preparation, since the caller has just named the shapes it is about to run.

Wouldn't it be more reasonable to do most of the preparation only in dispatch? With a call to computeShapes() the user does not promise to call dispatch afterwards. They might just to computeShapes to select to among different output shapes.

[theHamsta]

the post-compilation on dispatch for specialization might be finger printable. Also, the implementation hitting an very unfortunate shape might have a finger printable cost that allows guesses about the used HW and implementation.

[miaobin]

Wouldn't it be more reasonable to do most of the preparation only in dispatch?

Fair point — the wording was assuming something the API doesn't guarantee, and I've fixed that.
I would still keep computeShapes() as a natural place for preparation rather than moving all of it to dispatch. In the dominant pattern the call is a good predictor: a framework calls it precisely to size the output tensors for a dispatch it is about to make, since dispatch() takes caller-allocated outputs. Your case — comparing candidate shapes and running none of them — is real too, and an implementation cannot distinguish the two, so the section now frames this as a signal whose strength depends on why it was called, and notes that the call is synchronous, which bounds what can reasonably be done there.

The post-compilation on dispatch for specialization might be finger printable.

You're right, and timing is device information. I've reworked that opening to separate the two surfaces:

Declarative (unchanged): no value the API returns carries device or environment information a static graph does not.
Timing (widened): one graph dispatched over many shapes exposes a cost curve rather than the single build() duration, and its discontinuities can suggest the backend or the class of hardware.

I don't think there is a WebNN-level mitigation — constant-time execution across shapes isn't realistic for this API. Two things bound it: the work WebNN itself adds at dispatch is data-independent and cached per set of input shapes, so the variable part is the backend's own re-specialization, which the underlying runtime performs with or without WebNN; and explicit specialization would make that cost something the caller asked for rather than an implicit cost at dispatch.

[fdwr]

I would still keep computeShapes() as a natural place for preparation rather than moving all of it to dispatch.

👍 The more deferred work we squeeze into dispatch, the less dispatch behaves as originally intended - a lightweight execution of what has already been prepared. It's very common in graphics workloads (especially games) to want to frontload such blocking operations that would trigger memory planning and allocation. Unfortunately very few ML libraries nicely delineate between memory allocation and execution steps. WebNN at least offering (but not requiring) that step would afford implementations to the chance frontload things, reducing the stutter of the costly first frame.

[theHamsta]

At the moment, ONNX runtime would not do much shape analysis. ONNX, the library around the serialization format (e.g. onnx.shape_inference.infer_shapes_path), does some propagation of dim_param strings:

• it propagates the dim_param string to other intermediates and outputs where it can guarantee the equality of the values
• it usually generates new dim_params when an input dim_param is transformed in some way (e.g. a shape is padded)

PyTorch tries to do some symbolic shape inference via SymPy. Even more with dynamo, it would then even embed the symbolic restrictions into the ONNX shape strings. We could have similar capabilities in WebNN if we want and also guarantee that dertermining those expressions is always possible

WebNN would be free to impose some additional restrictions and validation on the input graph for which it will later require certain runtime behavior even if inference framework themselves don't perform that kind of analysis.

For onnxruntime with a TRT backend, TRT performs especially on the min,max,opt bounds symbolic shape inference and will report inconsistency errors on graph build. Especially, implementation that do actual kernel compilation (like some torch.compile backends) do symbolic inference because they can't afford to recompile on any shape change and want to have as much a-priori knowledge for compilation.

For the current proposal, would it be legitimate for an implementation to return errors from their symbolic analysis on MLGraphBuilder.build or would then need to defer their error reporting for computeShapes/dispatch?

script.py

[theHamsta]

To be clear I'm not necessarily against the current proposal. I'm also still undecided which of the approaches would be best for implementations or users to have.

As a ML devs, mismatching shapes in pytorch and ONNX just have always haunted me

[miaobin]

Would it be legitimate for an implementation to return errors from their symbolic analysis on MLGraphBuilder.build or would then need to defer their error reporting for computeShapes/dispatch?

I apologize that some of the descriptions caused confusion. I was likely preparing the explainer while simultaneously working on the POC, and I ended up conflating the implementation details with the specification. Because the section read as though the predicate was the rule. I've split the two apart.

The invariant is one-sided: an implementation must not reject a graph that some assignment of concrete input shapes would make valid. Nothing obliges it to be any good at detecting the rest. The three-valued predicate in the explainer is the weakest way to honor that — it only proves a contradiction between two static dimensions — so an implementation carrying symbolic expressions may reject strictly more, and earlier. Concatenating two [batch, seq] tensors along axis 1 doubles the element count, so reshaping that result back to [batch, seq] can never hold; our deferred-validation model sees only an unknown for the concatenated axis and defers, while a symbolic one knows the axis is 2 * seq, proves the graph invalid, and may fail build(). The one consequence worth stating is that a successful build() is then not a portable promise that a graph will ever run, only that this implementation could not prove otherwise — which I've now written down.

[theHamsta]

for those operators, symbolic shape inference would be inherently more difficult as the symbolic expression propagation would no longer be just shapes of tensors but also allowed to be actual data of the tensors.

The consequence of such operators is also that computeShape might be that it is possible to construct networks where computeShape is not simpler by any amount than actual dispatch:

• you pass you data as shapes
• you transform shape into data via the shape operator
• you run inference on the data (convolutions, matrix multiply)
• you transform data back into shape via those dynamic variants

So, your WebNN implementation is required to run actual inference during computeShapes and therefore also needs to have access to the weights. You could do that with the actual backend framework, but this framework will do its own kind of validation and inference behavior. How can a separate WebNN validation and shape inference logic from the validation and shape inference logic of my backend? I also know that many implementation will not be able to perform arbitrary meta on shape-to-data-and-back. Some stuff will work other things not. What does WebNN actually require implementation to work?

Another open question here is then for computeShape or shape inference in general, is it expected that the runtime on CPU or the accelerator runs the shape computation? For ONNX runtime, it is something weird in-between. TensorRT performs some analysis and tries to separate data paths where some tensors are shape tensors which ought to be interpreted by the runtime and others are GPU tensors which will be calculated by the runtime. Others, are kind of shape tensors, but they depend on inputs (computeShapes does not allow to depend on input data, but it still allows to transform input shapes into data) or are truely dispatch dependent (NonZero). Such tensors can just be interpreted by an accelerator by actually performing the calculation on a worst case maximum of memory and then output the calculated output side. A runtime has to classify those cases and it will work for most constructed network. However, there exist constructible networks that will just not work with most runtime. Which of those are required by WebNN to work?

When both computeShape and dispatch may use network constants where should I put them? On the accelerator or the runtime? Also this has to be determined by some kind of analysis. Some of those constant are required both for shape inference and inference. Would it make sense to separate those data paths and have one type MLOperand for data evaluated during dispatch and another MLShapeOperand which is only interpreted on CPU by the runtime? I can go from input shapes into calculations of MLShapeOperand which I interpret, but I have separate them for the data that is calculated during dispatch and on the accelerator. I would prefer to be able to reason tensor shapes. MLShapeOperand would be interpreted during computeShapes (or the validation of dispatch), MLOperand is interpreted from the accelerator during dispatch.

In RustNN, we currently construct the graph and translate it to the semantics of the backend framework. We preferably put the constant to the accelerator and the accelerator framework. We don't want to keep a copy of the constant for the RustNN runtime. If it comes now to dispatch or computeShape, we can either rely on semantics of the implementation framework to do computeShape/dispatch which might not agree with WebNN semantics (or whether it agrees is determined on build) or I duplicate all inference logic into something that just does shape inference conforming to WebNN.

[reillyeon]

I had a similar concern that computing shapes could effectively require executing the model but I convinced myself that the intent here is that the compute graph for shapes is separate from the compute graph for tensors. However the reuse of the MLOperand type makes that hard to see. We probably need to make this explicit by either adding an isShape property to an operand (forbidding non-shape operators from being used to compute shapes) or by using a completely different operand type.

[miaobin]

What does WebNN actually require implementation to work?

That was a real gap. The explainer said "no blessed list of operators" without saying what must therefore work, which reads as unbounded. It conflated two questions, and I've split them:

What a graph may express is not restricted by operator identity. Restricting it that way is what @fdwr objected to earlier, and I still think that's right. What is restricted is the chain's root: shape() outputs and build-time constants only. That restriction is the load-bearing one. It is why a shape chain cannot become the model.

ONNX answers the same question with a deliberately narrow closed set — the 11 operators carrying a data-propagation function (Add, Cast, Concat, Gather, Mul, Shape, Size, Slice, Squeeze, Sub, Unsqueeze) while ORT's symbolic_shape_infer.py covers a much wider set. A required minimum for WebNN belongs between the two. Naming it is work this proposal still owes, and I've written that down rather than leaving it implied; for calibration, the prototype today is a strict superset of the ONNX set.

Is it expected that the runtime on CPU or the accelerator runs the shape computation?

On CPU. A shape chain is integer work on vectors of at most rank length, and an implementation that cannot evaluate one fails cleanly rather than dispatching it.

Which of those are required by WebNN to work?

No classification is needed because the boundary is very clear: a chain bottoms out either at shape() and constants, or at a tensor's data. And that route isn't merely an implementation choice for us: there is no worst case to allocate (the model is deliberately unbounded, so bounded dimensions are a precondition), dispatch() takes caller-allocated outputs so the produced size would need a way back, and computeShapes() could not answer for such an output at all. I've expanded the Tensor-Derived open question to record that.

When both computeShape and dispatch may use network constants where should I put them?

That is a good question. Similar to the approach you described for RustNN, WebNN currently uploads constants to the service side as soon as they are created. Indeed, to perform validation based on actual input shapes, we need to retain certain constants related to shape inference. And the POC additionally caps each collected constant at a shape vector's worth of elements (16, covering Pad's 2 * max_rank). So you keep a handful of ≤16-element vectors, not anything model-sized, and the constants that are needed for both paths cost about a hundred bytes of duplication rather than a second copy of the model.
About the compute logic what we did is much smaller than "duplicate all inference logic". It is an integer interpreter over roughly twenty operators on vectors of at most rank length — about 880 lines in our implementation.

About MLShapeOperand

I added a section titled "Making the shape subgraph explicit" to the "Open Questions" section to facilitate further discussion.

[theHamsta]

That is a good question. Similar to the approach you described for RustNN, WebNN currently uploads constants to the service side as soon as they are created. Indeed, to perform validation based on actual input shapes, we need to retain certain constants related to shape inference.

Ah, so the intent is that computeShape is performed by the browser WebNN runtime on a certain sub-section of the graph and the rest of the graph is then performed by some accelerator framework. So, an implementation would try to split the graph into those two parts. Both for the shape graph and for the accelerator graph I need an implementation of WebNN to run all (or a fraction) the operators.

On CPU. A shape chain is integer work on vectors of at most rank length, and an implementation that cannot evaluate one fails cleanly rather than dispatching it.

Is there any guidance on how complete such a shape chain should be in comparison to a full WebNN implementation? Or would a typical approach be to just run it with the CPU backend of your WebNN implementation? Or just an ad-hoc interpreter for what you think that networks would typically use?

[theHamsta]

TRT has a function to get the output shapes given certain input shapes. ONNX runtime can only try to evaluate symbolic inference if it has been performed before on the ONNX, for the usual case, inference is required to know the output shapes as they also allow operators like NonZero where only execution determines the final output shape.

For the shape interpreter, I could create a new hand-written implementation of it or run ONNX runtime inference only on the shape graph part of the overall graph.

[miaobin]

To be precise about what the prototype does: we identify the shape subgraph by walking backward from the shape operands of the *Dynamic operators. That walk bottoms out at shape() outputs and build-time constants — it cannot bottom out anywhere else, because the interpreter has no access to tensor data (an input
operand simply evaluates to "unknown"). That root restriction is what keeps the shape subgraph a small subgraph rather than the whole model.

A standalone interpreter evaluates that subgraph on the CPU. What it implements in the POC is the set (Add, Cast, Concat, Gather, Mul, Shape, Size, Slice, Squeeze, Sub, Unsqueeze): arithmetic and structural transforms over small 1-D vectors, plus enough floating-point for cases like floor(dim * scale). It grew from what the models we tested actually needed rather than from a designed list; [symbolic_shape_infer.py](https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/symbolic_shape_infer.py#L130) is the wider reference we point at when asking where a required floor should sit.

Worth stressing: this does not partition the graph or move any work off the accelerator. The graph handed to the backend is unchanged — shape() and the *Dynamic operators are emitted as ordinary nodes, and the backend executes the whole graph however it sees fit. The CPU evaluation is a side computation, used for validation and for answering computeShapes().

Two things make it necessary:

1. Allocation ordering. dispatch() takes caller-allocated output MLTensors, and requires each one's descriptor to equal the graph's output descriptor — same data type, same shape, dimension by dimension. An MLTensor's shape is fixed at creation, and dispatch() returns undefined synchronously, so it has no way to hand back tensors it allocated itself. The caller — an application, or a framework such as ORT Web — therefore has to know the output shapes for the given input shapes before it calls. Without computeShapes() it would have to reimplement WebNN's shape inference just to size those allocations.

2. Validation has to happen before work reaches the backend. The renderer is untrusted, so the privileged service re-runs the same inference independently rather than taking the renderer's word for it. And dispatch() has no error-reporting channel: a failure that reaches the backend cannot be reported per-dispatch, only escalated to context loss — in our implementation, a session-run failure currently tears down every WebNN context in the GPU process. So "can these shapes be resolved, and do they match what the caller allocated?" has to be answered up front, on the CPU, before anything is submitted.

[theHamsta]

MLGraph.computeShapes is quite similar to MLContext.dispatch. Though, ccomputeShapes is on graph and dispatch is on MLContext.

[miaobin]

dispatch() binds MLTensors on the context, so it belongs to MLContext and takes the graph as an argument.
computeShapes() touches neither tensors nor the device. It reads the graph's topology and the caller's integer shapes. So it sits on MLGraph, and that is also what lets it be synchronous.

[theHamsta]

Regarding the separations of constants: we could separate MLGraphBuilder.constant from MLGraphBuilder.shape_constant

• computeShape would only be allowed to read from shape_constant but not from a regular constant
• regular inference would be allowed to read both

This would allow to immediately send .constant to the accelerator without needing to wait for the user to construct the full graph and be able to analyze it.

To me it seems that actually computeShape never needs to read the contents of constants unless they are input to one of the rank changing or dynamic operator variants. Those operators basically transform the input chain to their shape argument from the dispatch part of the graph to the computeShape part of the computation.

[theHamsta]

Please feel free to close my comments! If you feel that there is nothing actionable. Most of the comments or mainly to clarify our understanding of this proposal. Please excuse if this generates any kind of noise for you!

[miaobin]

Not noise at all! I really appreciate the thorough read. A number of your comments led directly to changes in the explainer, and your comments are genuinely useful for pressure-testing whether the explanation actually holds up. Thanks for taking the time!