onnxruntime
2650ca13 - Fix JSEP Reshape and Unsqueeze schema registrations (#33749)

Commit
7 days ago
Fix JSEP Reshape and Unsqueeze schema registrations (#33749) ### Description <!-- Describe your changes. --> Fix JSEP kernel selection for existing `Reshape` and `Unsqueeze` operations in newer models, with regression coverage through opset 28. - Replace the schema-21-only registrations with bounded schema 21-24 registrations, add current schema-25 registrations, and update both the declarations and registry entries. - Keep `JsepSupportedDataTypes()` unchanged: FLOAT, FLOAT16, INT32, and UINT32. Shape/axes inputs remain INT64 in CPU memory, and input/output aliasing is unchanged. - Add browser regressions that create real WebGPU sessions with graph optimizations and CPU EP fallback disabled. Cover every model opset from 21 through 28, all four supported types, initializer and dynamic shape/axes inputs, negative axes, Reshape zero/inferred dimensions and `allowzero=1`, and rejection of unsupported FLOAT64/FLOAT4E2M1/FLOAT8E8M0/UINT2/INT2 types. - Regenerate the JSEP operator coverage table. This is a bounded correctness fix for existing JSEP view operations, not support for new data types or a general opset update. Native WebGPU already has explicit registrations for these schemas; it does not need the same registration fix. Other operators and graph capture are outside this change. ### Motivation and Context <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. --> Valid FLOAT32 models fail GPU-only session creation with `onnxruntime-web/all` 1.31.1 (JSEP). Both operators succeed at opset 22 and fail at opsets 23 and 24 on Intel WebGPU, despite passing full ONNX checking and native CPU numerical checks: ```text webgpu kernel not found in registries for Op type: Reshape This session contains graph nodes that are assigned to the default CPU EP, but fallback to CPU EP has been explicitly disabled by the user. ``` `KernelRegistry::VerifyVersion()` matches the node's operator schema version, not the model's imported opset. An unbounded registration deliberately requires an exact schema-version match; only explicitly finite ranges match multiple schema versions. Consequently, a registration starting at 21 does not match schemas 23, 24, or 25. For both operators, the pinned ONNX dependency resolves these model opsets as follows: | Model opset | Operator schema | Registration after this fix | |---|---|---| | 21, 22 | 21 | 21-24 | | 23 | 23 | 21-24 | | 24 | 24 | 21-24 | | 25, 26, 27, 28 | 25 | 25 | Thus, there is no model-opset-24 cap. The current schema-25 registration also applies to later model opsets that continue to select schema 25; it does not blindly promise compatibility with future revised operator schemas. The pinned ONNX schemas retain the same valid-input shape/axes and `allowzero` semantics for JSEP's existing types: schema 23 adds FLOAT4E2M1 to the allowed types, schema 24 adds FLOAT8E8M0, and schema 25 adds UINT2/INT2 through IR 13. None of those newly allowed types is advertised by this fix. Reshape's schema-25 shape inference also adds overflow checks, which run in ONNX and require no JSEP kernel change. Minimal reproduction (IR 10, opset 24): | Operator | Input | INT64 initializer | Expected output | |---|---|---|---| | Reshape | FLOAT `[2,2]`, values `[1,2,3,4]` | shape `[4]` | FLOAT `[4]`, same values | | Unsqueeze | FLOAT `[4]`, values `[1,2,3,4]` | axes `[0]` | FLOAT `[1,4]`, same values | Load either model in a browser using `/all`: ```js import * as ort from 'onnxruntime-web/all'; const session = await ort.InferenceSession.create(modelBytes, { executionProviders: ['webgpu'], graphOptimizationLevel: 'disabled', extra: { session: { disable_cpu_ep_fallback: '1' } }, }); const result = await session.run({ x: new ort.Tensor('float32', Float32Array.of(1, 2, 3, 4), inputDims), }); ``` The added tests construct these models with the repository's existing ONNX protobuf builder and exercise session creation, provider assignment, and execution rather than calling a JavaScript kernel directly. --------- Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Author
Parents
Loading