sync 91d990483a17

Browse files

Files changed (8) hide show

README.md +15 -11
build/webgpu/bench.json +0 -1
build/webgpu/datamove-elementwise-copy.wgsl.jinja +0 -3
build/webgpu/manifest.json +229 -470
build/webgpu/metadata.json +24 -9
build/webgpu/summean-broadcast.wgsl.jinja +13 -16
build/webgpu/summean-vec4.wgsl.jinja +4 -7
build/webgpu/test.json +6 -7

README.md CHANGED Viewed

@@ -18,19 +18,19 @@ See the [ONNX `Mean` spec](https://onnx.ai/onnx/operators/onnx__Mean.html) for t
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `A` | `a` | `T` | — | — | First input tensor. | required |
-| `B` | `b` | `T` | — | — | Second input tensor, broadcast-compatible with A. | optional |
-| `C` | `c` | `T` | — | — | Third input tensor, broadcast-compatible with A and B. | optional |
-| `D` | `d` | `T` | — | — | Fourth input tensor, broadcast-compatible with A, B, and C. | optional |
-| `E` | `e` | `T` | — | — | Fifth input tensor, broadcast-compatible with all other inputs. | optional |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `mean` | `y` | `T` | derived | derived; see description | Elementwise mean of all provided input tensors. | required |
 ## Type constraints
@@ -40,7 +40,7 @@ See the [ONNX `Mean` spec](https://onnx.ai/onnx/operators/onnx__Mean.html) for t
 ## Files
-- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
@@ -50,10 +50,14 @@ See the [ONNX `Mean` spec](https://onnx.ai/onnx/operators/onnx__Mean.html) for t
 ## Use with `@huggingface/kernels`
-The loader derives every required output's shape and logical dtype from the manifest contract and this call.
-It then allocates the result tensors automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `a` | `A` | `T` | — | — | First input tensor. | required |
+| `b` | `B` | `T` | — | — | Second input tensor, broadcast-compatible with A. | optional |
+| `c` | `C` | `T` | — | — | Third input tensor, broadcast-compatible with A and B. | optional |
+| `d` | `D` | `T` | — | — | Fourth input tensor, broadcast-compatible with A, B, and C. | optional |
+| `e` | `E` | `T` | — | — | Fifth input tensor, broadcast-compatible with all other inputs. | optional |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `y` | `mean` | `T` | derived | derived | Elementwise mean of all provided input tensors. | required |
 ## Type constraints
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.Mean",
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
   "cases": [
     {

 {
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
   "cases": [
     {

build/webgpu/datamove-elementwise-copy.wgsl.jinja CHANGED Viewed

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;





1	{{ env.wgsl.resourceDeclarations }}
2
3	const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;

build/webgpu/manifest.json CHANGED Viewed

@@ -2,432 +2,181 @@
   "domain": "ai.onnx",
   "name": "Mean",
   "sinceVersion": 13,
-  "description": "Computes the elementwise mean of one or more input tensors with multidirectional (NumPy-style) broadcasting. All inputs and the output share the same data type.",
-  "inputs": [
-    { "role": "A", "dtype": "T", "description": "First input tensor." },
-    { "role": "B", "dtype": "T", "optional": true, "description": "Second input tensor, broadcast-compatible with A." },
-    {
-      "role": "C",
-      "dtype": "T",
-      "optional": true,
-      "description": "Third input tensor, broadcast-compatible with A and B."
-    },
-    {
-      "role": "D",
-      "dtype": "T",
-      "optional": true,
-      "description": "Fourth input tensor, broadcast-compatible with A, B, and C."
-    },
-    {
-      "role": "E",
-      "dtype": "T",
-      "optional": true,
-      "description": "Fifth input tensor, broadcast-compatible with all other inputs."
-    }
-  ],
-  "outputs": [
-    {
-      "role": "mean",
       "dtype": "T",
-      "rank": "max(ranks.A, ranks.B if present.b else 0, ranks.C if present.b and present.c else 0, ranks.D if present.b and present.c and present.d else 0, ranks.E if present.b and present.c and present.d and present.e else 0)",
-      "shape": "variadicShape",
-      "description": "Elementwise mean of all provided input tensors."
     }
-  ],
-  "typeConstraints": { "T": ["float32", "float16"] },
-  "args": {
-    "a": { "kind": "tensor", "semantic": "A", "role": "input" },
-    "b": { "kind": "tensor", "semantic": "B", "role": "input", "required": false },
-    "c": { "kind": "tensor", "semantic": "C", "role": "input2", "required": false },
-    "y": { "kind": "tensor", "semantic": "mean", "role": "output" },
-    "d": { "kind": "tensor", "semantic": "D", "role": "input3", "required": false },
-    "e": { "kind": "tensor", "semantic": "E", "role": "input4", "required": false }
   },
-  "tunables": { "WORKGROUP_SIZE": 256 },
   "derive": {
     "variadicInputCount": "1 + (1 if present.b else 0) + (1 if (present.b and present.c) else 0) + (1 if (present.b and present.c and present.d) else 0) + (1 if (present.b and present.c and present.d and present.e) else 0)",
-    "variadicShape": "broadcastShape(broadcastShape(broadcastShape(broadcastShape(shapes.A, shapes.B), shapes.C), shapes.D), shapes.E) if present.b and present.c and present.d and present.e else (broadcastShape(broadcastShape(broadcastShape(shapes.A, shapes.B), shapes.C), shapes.D) if present.b and present.c and present.d else (broadcastShape(broadcastShape(shapes.A, shapes.B), shapes.C) if present.b and present.c else (broadcastShape(shapes.A, shapes.B) if present.b else shapes.A)))",
     "flatVec4OutputOk": "numel(shapes.y) > 0 and numel(shapes.y) % 4 == 0 and f16Ok(dtypes.T)",
     "broadcastOutputOk": "f16Ok(dtypes.T)",
     "broadcastVec4OutputOk": "ranks.y >= 1 and dim(shapes.y, ranks.y - 1) > 0 and dim(shapes.y, ranks.y - 1) % 4 == 0 and f16Ok(dtypes.T)",
-    "aBroadcastVec4Ok": "ranks.A <= ranks.y and (ranks.A == 0 or dim(shapes.A, ranks.A - 1) == 1 or dim(shapes.A, ranks.A - 1) == dim(shapes.y, ranks.y - 1))",
-    "bBroadcastVec4Ok": "not present.b or (ranks.B <= ranks.y and (ranks.B == 0 or dim(shapes.B, ranks.B - 1) == 1 or dim(shapes.B, ranks.B - 1) == dim(shapes.y, ranks.y - 1)))",
-    "cBroadcastVec4Ok": "not present.c or (ranks.C <= ranks.y and (ranks.C == 0 or dim(shapes.C, ranks.C - 1) == 1 or dim(shapes.C, ranks.C - 1) == dim(shapes.y, ranks.y - 1)))",
-    "dBroadcastVec4Ok": "not present.d or (ranks.D <= ranks.y and (ranks.D == 0 or dim(shapes.D, ranks.D - 1) == 1 or dim(shapes.D, ranks.D - 1) == dim(shapes.y, ranks.y - 1)))"
   },
-  "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
-  "bindingSets": {
-    "identity": [
-      { "name": "x", "arg": "a", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "y", "arg": "y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] }
-      }
-    ],
-    "vec4Two": [
-      {
-        "name": "a",
-        "arg": "a",
-        "semantic": "A",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "b",
-        "arg": "b",
-        "semantic": "B",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      { "name": "y", "arg": "y", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }] }
-      }
-    ],
-    "vec4Three": [
-      {
-        "name": "a",
-        "arg": "a",
-        "semantic": "A",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "b",
-        "arg": "b",
-        "semantic": "B",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "c",
-        "arg": "c",
-        "semantic": "C",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      { "name": "y", "arg": "y", "semantic": "mean", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }] }
-      }
-    ],
-    "vec4Four": [
-      {
-        "name": "a",
-        "arg": "a",
-        "semantic": "A",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "b",
-        "arg": "b",
-        "semantic": "B",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "c",
-        "arg": "c",
-        "semantic": "C",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "d",
-        "arg": "d",
-        "semantic": "D",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      { "name": "y", "arg": "y", "semantic": "mean", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }] }
-      }
-    ],
-    "scalarTwo": [
-      { "name": "a", "arg": "a", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "b", "arg": "b", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "y", "arg": "y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] }
-      }
-    ],
-    "scalarThree": [
-      { "name": "a", "arg": "a", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "b", "arg": "b", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "c", "arg": "c", "semantic": "C", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "y", "arg": "y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] }
-      }
-    ],
-    "scalarFour": [
-      { "name": "a", "arg": "a", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "b", "arg": "b", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "c", "arg": "c", "semantic": "C", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "d", "arg": "d", "semantic": "D", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "y", "arg": "y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] }
-      }
-    ],
-    "vec4Five": [
-      {
-        "name": "a",
-        "arg": "a",
-        "semantic": "A",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "b",
-        "arg": "b",
-        "semantic": "B",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "c",
-        "arg": "c",
-        "semantic": "C",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "d",
-        "arg": "d",
-        "semantic": "D",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "e",
-        "arg": "e",
-        "semantic": "E",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      { "name": "y", "arg": "y", "semantic": "mean", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }] }
-      }
-    ],
-    "scalarFive": [
-      { "name": "a", "arg": "a", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "b", "arg": "b", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "c", "arg": "c", "semantic": "C", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "d", "arg": "d", "semantic": "D", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "e", "arg": "e", "semantic": "E", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "y", "arg": "y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] }
-      }
-    ],
-    "vec4BroadcastThree": [
-      {
-        "name": "a",
-        "arg": "a",
-        "semantic": "A",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aElement"
-      },
-      {
-        "name": "b",
-        "arg": "b",
-        "semantic": "B",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$bElement"
-      },
-      {
-        "name": "c",
-        "arg": "c",
-        "semantic": "C",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$cElement"
-      },
-      { "name": "y", "arg": "y", "semantic": "mean", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }] }
-      }
-    ],
-    "vec4BroadcastFour": [
-      {
-        "name": "a",
-        "arg": "a",
-        "semantic": "A",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aElement"
-      },
-      {
-        "name": "b",
-        "arg": "b",
-        "semantic": "B",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$bElement"
-      },
-      {
-        "name": "c",
-        "arg": "c",
-        "semantic": "C",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$cElement"
-      },
-      {
-        "name": "d",
-        "arg": "d",
-        "semantic": "D",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$dElement"
-      },
-      { "name": "y", "arg": "y", "semantic": "mean", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }] }
-      }
-    ]
   },
   "variants": [
     {
       "id": "single_input_identity",
       "priority": 30,
-      "when": ["variadicInputCount == 1", "ranks.A == ranks.y", "numel(shapes.A) == numel(shapes.y)", "f16Ok(dtypes.T)"],
       "passes": [
         {
           "id": "main",
           "name": "Mean",
           "shader": "datamove-elementwise-copy.wgsl.jinja",
-          "bindings": "identity",
-          "dispatch": { "gridStride": "numel(shapes.y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "same_shape_vec4_two_input",
       "priority": 20,
-      "when": ["variadicInputCount == 2", "sameShape(shapes.A, shapes.y)", "sameShape(shapes.B, shapes.y)", "flatVec4OutputOk"],
-      "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Mean.vec4",
-          "source": { "shader": "summean-vec4.wgsl.jinja", "inputs": { "op": "\"mean\"", "hasC": "false" } },
-          "bindings": "vec4Two",
-          "dispatch": { "threads": "numel(shapes.y) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "same_shape_vec4_three_input",
       "priority": 25,
-      "when": ["variadicInputCount == 3", "sameShape(shapes.A, shapes.y)", "sameShape(shapes.B, shapes.y)", "sameShape(shapes.C, shapes.y)", "flatVec4OutputOk"],
-      "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Mean.vec4_3",
-          "source": { "shader": "summean-vec4.wgsl.jinja", "inputs": { "op": "\"mean\"", "hasC": "true" } },
-          "bindings": "vec4Three",
-          "dispatch": { "threads": "numel(shapes.y) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "same_shape_vec4_four_input",
       "priority": 27,
-      "when": ["variadicInputCount == 4", "sameShape(shapes.A, shapes.y)", "sameShape(shapes.B, shapes.y)", "sameShape(shapes.C, shapes.y)", "sameShape(shapes.D, shapes.y)", "flatVec4OutputOk"],
-      "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Mean.vec4_3",
-          "source": {
-            "shader": "summean-vec4.wgsl.jinja",
-            "inputs": { "op": "\"mean\"", "hasC": "true", "hasD": "true" }
-          },
-          "bindings": "vec4Four",
-          "dispatch": { "threads": "numel(shapes.y) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "same_shape_vec4_five_input",
       "priority": 28,
-      "when": ["variadicInputCount == 5", "sameShape(shapes.A, shapes.y)", "sameShape(shapes.B, shapes.y)", "sameShape(shapes.C, shapes.y)", "sameShape(shapes.D, shapes.y)", "sameShape(shapes.E, shapes.y)", "flatVec4OutputOk"],
-      "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Mean.vec4_35",
-          "source": {
-            "shader": "summean-vec4.wgsl.jinja",
-            "inputs": {
-              "op": "\"mean\"",
-              "hasC": "true",
-              "hasD": "true",
-              "hasE": "true",
-              "extraInputs": "[\"c\", \"d\", \"e\"]"
-            }
           },
-          "bindings": "vec4Five",
-          "dispatch": { "threads": "numel(shapes.y) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "broadcast_two_input",
-      "when": ["variadicInputCount == 2", "ranks.A <= ranks.y", "ranks.B <= ranks.y", "broadcastOutputOk"],
       "passes": [
         {
           "id": "main",
           "name": "Mean",
-          "source": {
-            "shader": "summean-broadcast.wgsl.jinja",
-            "inputs": {
-              "aShape": "shapes.A",
-              "bShape": "shapes.B",
-              "yShape": "shapes.y",
-              "aRank": "ranks.A",
-              "bRank": "ranks.B",
-              "yRank": "ranks.y",
-              "hasC": "false",
-              "op": "\"mean\""
-            }
           },
-          "bindings": "scalarTwo",
-          "dispatch": { "threads": "numel(shapes.y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
@@ -435,65 +184,69 @@
       "id": "broadcast_three_input_vec4",
       "priority": 24,
       "when": ["variadicInputCount == 3", "broadcastVec4OutputOk", "aBroadcastVec4Ok", "bBroadcastVec4Ok", "cBroadcastVec4Ok"],
-      "constants": {
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
-        "aElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if ranks.A > 0 and dim(shapes.A, ranks.A - 1) != 1 else dtypes.T",
-        "bElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if ranks.B > 0 and dim(shapes.B, ranks.B - 1) != 1 else dtypes.T",
-        "cElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if ranks.C > 0 and dim(shapes.C, ranks.C - 1) != 1 else dtypes.T"
       },
       "passes": [
         {
           "id": "main",
           "name": "Mean.Broadcast3Vec4",
-          "source": {
-            "shader": "summean-broadcast.wgsl.jinja",
-            "inputs": {
-              "aShape": "shapes.A",
-              "bShape": "shapes.B",
-              "cShape": "shapes.C",
-              "yShape": "shapes.y",
-              "aRank": "ranks.A",
-              "bRank": "ranks.B",
-              "cRank": "ranks.C",
-              "yRank": "ranks.y",
-              "aVector": "ranks.A > 0 and dim(shapes.A, ranks.A - 1) != 1",
-              "bVector": "ranks.B > 0 and dim(shapes.B, ranks.B - 1) != 1",
-              "cVector": "ranks.C > 0 and dim(shapes.C, ranks.C - 1) != 1",
-              "hasC": true,
-              "vectorized": true,
-              "op": "\"mean\""
-            }
           },
-          "bindings": "vec4BroadcastThree",
-          "dispatch": { "threads": "numel(shapes.y) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "broadcast_three_input",
       "priority": 20,
-      "when": ["variadicInputCount == 3", "ranks.A <= ranks.y", "ranks.B <= ranks.y", "ranks.C <= ranks.y", "broadcastOutputOk"],
       "passes": [
         {
           "id": "main",
           "name": "Mean",
-          "source": {
-            "shader": "summean-broadcast.wgsl.jinja",
-            "inputs": {
-              "aShape": "shapes.A",
-              "bShape": "shapes.B",
-              "cShape": "shapes.C",
-              "yShape": "shapes.y",
-              "aRank": "ranks.A",
-              "bRank": "ranks.B",
-              "cRank": "ranks.C",
-              "yRank": "ranks.y",
-              "hasC": "true",
-              "op": "\"mean\""
-            }
           },
-          "bindings": "scalarThree",
-          "dispatch": { "threads": "numel(shapes.y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
@@ -501,108 +254,114 @@
       "id": "broadcast_four_input_vec4",
       "priority": 26,
       "when": ["variadicInputCount == 4", "broadcastVec4OutputOk", "aBroadcastVec4Ok", "bBroadcastVec4Ok", "cBroadcastVec4Ok", "dBroadcastVec4Ok"],
-      "constants": {
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
-        "aElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if ranks.A > 0 and dim(shapes.A, ranks.A - 1) != 1 else dtypes.T",
-        "bElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if ranks.B > 0 and dim(shapes.B, ranks.B - 1) != 1 else dtypes.T",
-        "cElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if ranks.C > 0 and dim(shapes.C, ranks.C - 1) != 1 else dtypes.T",
-        "dElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if ranks.D > 0 and dim(shapes.D, ranks.D - 1) != 1 else dtypes.T"
       },
       "passes": [
         {
           "id": "main",
           "name": "Mean.Broadcast4Vec4",
-          "source": {
-            "shader": "summean-broadcast.wgsl.jinja",
-            "inputs": {
-              "aShape": "shapes.A",
-              "bShape": "shapes.B",
-              "cShape": "shapes.C",
-              "dShape": "shapes.D",
-              "yShape": "shapes.y",
-              "aRank": "ranks.A",
-              "bRank": "ranks.B",
-              "cRank": "ranks.C",
-              "dRank": "ranks.D",
-              "yRank": "ranks.y",
-              "aVector": "ranks.A > 0 and dim(shapes.A, ranks.A - 1) != 1",
-              "bVector": "ranks.B > 0 and dim(shapes.B, ranks.B - 1) != 1",
-              "cVector": "ranks.C > 0 and dim(shapes.C, ranks.C - 1) != 1",
-              "dVector": "ranks.D > 0 and dim(shapes.D, ranks.D - 1) != 1",
-              "hasC": true,
-              "hasD": true,
-              "vectorized": true,
-              "op": "\"mean\""
-            }
           },
-          "bindings": "vec4BroadcastFour",
-          "dispatch": { "threads": "numel(shapes.y) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "broadcast_four_input",
       "priority": 22,
-      "when": ["variadicInputCount == 4", "ranks.A <= ranks.y", "ranks.B <= ranks.y", "ranks.C <= ranks.y", "ranks.D <= ranks.y", "broadcastOutputOk"],
       "passes": [
         {
           "id": "main",
           "name": "Mean",
-          "source": {
-            "shader": "summean-broadcast.wgsl.jinja",
-            "inputs": {
-              "aShape": "shapes.A",
-              "bShape": "shapes.B",
-              "cShape": "shapes.C",
-              "yShape": "shapes.y",
-              "aRank": "ranks.A",
-              "bRank": "ranks.B",
-              "cRank": "ranks.C",
-              "yRank": "ranks.y",
-              "hasC": "true",
-              "op": "\"mean\"",
-              "dShape": "shapes.D",
-              "dRank": "ranks.D",
-              "hasD": "true"
-            }
           },
-          "bindings": "scalarFour",
-          "dispatch": { "threads": "numel(shapes.y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "broadcast_five_input",
       "priority": 23,
-      "when": ["variadicInputCount == 5", "ranks.A <= ranks.y", "ranks.B <= ranks.y", "ranks.C <= ranks.y", "ranks.D <= ranks.y", "ranks.E <= ranks.y", "broadcastOutputOk"],
       "passes": [
         {
           "id": "main",
           "name": "Mean5",
-          "source": {
-            "shader": "summean-broadcast.wgsl.jinja",
-            "inputs": {
-              "aShape": "shapes.A",
-              "bShape": "shapes.B",
-              "cShape": "shapes.C",
-              "yShape": "shapes.y",
-              "aRank": "ranks.A",
-              "bRank": "ranks.B",
-              "cRank": "ranks.C",
-              "yRank": "ranks.y",
-              "hasC": "true",
-              "op": "\"mean\"",
-              "dShape": "shapes.D",
-              "dRank": "ranks.D",
-              "hasD": "true",
-              "hasE": "true",
-              "extraInputs": "[\"c\", \"d\", \"e\"]",
-              "eShape": "shapes.E",
-              "eRank": "ranks.E"
-            }
           },
-          "bindings": "scalarFive",
-          "dispatch": { "threads": "numel(shapes.y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     }

   "domain": "ai.onnx",
   "name": "Mean",
   "sinceVersion": 13,
+  "inputs": {
+    "a": { "onnx": "A", "dtype": "T" },
+    "b": { "onnx": "B", "dtype": "T", "optional": true },
+    "c": { "onnx": "C", "dtype": "T", "optional": true },
+    "d": { "onnx": "D", "dtype": "T", "optional": true },
+    "e": { "onnx": "E", "dtype": "T", "optional": true }
+  },
+  "outputs": {
+    "y": {
+      "onnx": "mean",
       "dtype": "T",
+      "rank": "max(ranks.a, ranks.b if present.b else 0, ranks.c if present.b and present.c else 0, ranks.d if present.b and present.c and present.d else 0, ranks.e if present.b and present.c and present.d and present.e else 0)",
+      "shape": "variadicShape"
     }
   },
+  "typeConstraints": { "T": ["float32", "float16"] },
+  "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
   "derive": {
     "variadicInputCount": "1 + (1 if present.b else 0) + (1 if (present.b and present.c) else 0) + (1 if (present.b and present.c and present.d) else 0) + (1 if (present.b and present.c and present.d and present.e) else 0)",
+    "variadicShape": "broadcastShape(broadcastShape(broadcastShape(broadcastShape(shapes.a, shapes.b), shapes.c), shapes.d), shapes.e) if present.b and present.c and present.d and present.e else (broadcastShape(broadcastShape(broadcastShape(shapes.a, shapes.b), shapes.c), shapes.d) if present.b and present.c and present.d else (broadcastShape(broadcastShape(shapes.a, shapes.b), shapes.c) if present.b and present.c else (broadcastShape(shapes.a, shapes.b) if present.b else shapes.a)))",
     "flatVec4OutputOk": "numel(shapes.y) > 0 and numel(shapes.y) % 4 == 0 and f16Ok(dtypes.T)",
     "broadcastOutputOk": "f16Ok(dtypes.T)",
     "broadcastVec4OutputOk": "ranks.y >= 1 and dim(shapes.y, ranks.y - 1) > 0 and dim(shapes.y, ranks.y - 1) % 4 == 0 and f16Ok(dtypes.T)",
+    "aBroadcastVec4Ok": "ranks.a <= ranks.y and (ranks.a == 0 or dim(shapes.a, ranks.a - 1) == 1 or dim(shapes.a, ranks.a - 1) == dim(shapes.y, ranks.y - 1))",
+    "bBroadcastVec4Ok": "not present.b or (ranks.b <= ranks.y and (ranks.b == 0 or dim(shapes.b, ranks.b - 1) == 1 or dim(shapes.b, ranks.b - 1) == dim(shapes.y, ranks.y - 1)))",
+    "cBroadcastVec4Ok": "not present.c or (ranks.c <= ranks.y and (ranks.c == 0 or dim(shapes.c, ranks.c - 1) == 1 or dim(shapes.c, ranks.c - 1) == dim(shapes.y, ranks.y - 1)))",
+    "dBroadcastVec4Ok": "not present.d or (ranks.d <= ranks.y and (ranks.d == 0 or dim(shapes.d, ranks.d - 1) == 1 or dim(shapes.d, ranks.d - 1) == dim(shapes.y, ranks.y - 1)))",
+    "scalar": "dtypes.T",
+    "usesF16": "dtypes.T == \"f16\""
   },
+  "bindings": {
+    "params": { "buffer": "uniform", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] },
+    "a": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
+    "b": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
+    "y_2": { "name": "y", "buffer": "storage", "elementType": "$vectorScalar" },
+    "params_2": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }]
+    },
+    "c": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
+    "d": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
+    "a_3": { "name": "a", "buffer": "read-only-storage", "elementType": "$aElement" },
+    "b_3": { "name": "b", "buffer": "read-only-storage", "elementType": "$bElement" },
+    "c_2": { "name": "c", "buffer": "read-only-storage", "elementType": "$cElement" }
   },
   "variants": [
     {
       "id": "single_input_identity",
       "priority": 30,
+      "when": ["variadicInputCount == 1", "ranks.a == ranks.y", "numel(shapes.a) == numel(shapes.y)", "f16Ok(dtypes.T)"],
       "passes": [
         {
           "id": "main",
           "name": "Mean",
           "shader": "datamove-elementwise-copy.wgsl.jinja",
+          "bindings": [{ "arg": "a", "name": "x", "elementType": "$scalar" }, "y", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "same_shape_vec4_two_input",
       "priority": 20,
+      "when": ["variadicInputCount == 2", "sameShape(shapes.a, shapes.y)", "sameShape(shapes.b, shapes.y)", "flatVec4OutputOk"],
+      "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Mean.vec4",
+          "shader": "summean-vec4.wgsl.jinja",
+          "derive": { "op": "\"mean\"", "hasC": "false" },
+          "bindings": ["a", "b", "y_2", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "same_shape_vec4_three_input",
       "priority": 25,
+      "when": ["variadicInputCount == 3", "sameShape(shapes.a, shapes.y)", "sameShape(shapes.b, shapes.y)", "sameShape(shapes.c, shapes.y)", "flatVec4OutputOk"],
+      "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Mean.vec4_3",
+          "shader": "summean-vec4.wgsl.jinja",
+          "derive": { "op": "\"mean\"", "hasC": "true" },
+          "bindings": ["a", "b", "c", "y_2", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "same_shape_vec4_four_input",
       "priority": 27,
+      "when": ["variadicInputCount == 4", "sameShape(shapes.a, shapes.y)", "sameShape(shapes.b, shapes.y)", "sameShape(shapes.c, shapes.y)", "sameShape(shapes.d, shapes.y)", "flatVec4OutputOk"],
+      "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Mean.vec4_3",
+          "shader": "summean-vec4.wgsl.jinja",
+          "derive": { "op": "\"mean\"", "hasC": "true", "hasD": "true" },
+          "bindings": ["a", "b", "c", "d", "y_2", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "same_shape_vec4_five_input",
       "priority": 28,
+      "when": ["variadicInputCount == 5", "sameShape(shapes.a, shapes.y)", "sameShape(shapes.b, shapes.y)", "sameShape(shapes.c, shapes.y)", "sameShape(shapes.d, shapes.y)", "sameShape(shapes.e, shapes.y)", "flatVec4OutputOk"],
+      "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Mean.vec4_35",
+          "shader": "summean-vec4.wgsl.jinja",
+          "derive": {
+            "op": "\"mean\"",
+            "hasC": "true",
+            "hasD": "true",
+            "hasE": "true",
+            "extraInputs": "[\"c\", \"d\", \"e\"]"
           },
+          "bindings": ["a", "b", "c", "d", { "arg": "e", "elementType": "$vectorScalar" }, "y_2", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "broadcast_two_input",
+      "when": ["variadicInputCount == 2", "ranks.a <= ranks.y", "ranks.b <= ranks.y", "broadcastOutputOk"],
       "passes": [
         {
           "id": "main",
           "name": "Mean",
+          "shader": "summean-broadcast.wgsl.jinja",
+          "derive": {
+            "aShape": "shapes.a",
+            "bShape": "shapes.b",
+            "yShape": "shapes.y",
+            "aRank": "ranks.a",
+            "bRank": "ranks.b",
+            "yRank": "ranks.y",
+            "hasC": "false",
+            "op": "\"mean\""
           },
+          "bindings": [{ "arg": "a" }, { "arg": "b" }, "y", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
       "id": "broadcast_three_input_vec4",
       "priority": 24,
       "when": ["variadicInputCount == 3", "broadcastVec4OutputOk", "aBroadcastVec4Ok", "bBroadcastVec4Ok", "cBroadcastVec4Ok"],
+      "derive": {
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "aElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if ranks.a > 0 and dim(shapes.a, ranks.a - 1) != 1 else dtypes.T",
+        "bElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if ranks.b > 0 and dim(shapes.b, ranks.b - 1) != 1 else dtypes.T",
+        "cElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if ranks.c > 0 and dim(shapes.c, ranks.c - 1) != 1 else dtypes.T"
       },
       "passes": [
         {
           "id": "main",
           "name": "Mean.Broadcast3Vec4",
+          "shader": "summean-broadcast.wgsl.jinja",
+          "derive": {
+            "aShape": "shapes.a",
+            "bShape": "shapes.b",
+            "cShape": "shapes.c",
+            "yShape": "shapes.y",
+            "aRank": "ranks.a",
+            "bRank": "ranks.b",
+            "cRank": "ranks.c",
+            "yRank": "ranks.y",
+            "aVector": "ranks.a > 0 and dim(shapes.a, ranks.a - 1) != 1",
+            "bVector": "ranks.b > 0 and dim(shapes.b, ranks.b - 1) != 1",
+            "cVector": "ranks.c > 0 and dim(shapes.c, ranks.c - 1) != 1",
+            "hasC": true,
+            "vectorizedSpec": true,
+            "op": "\"mean\""
           },
+          "bindings": ["a_3", "b_3", "c_2", "y_2", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "broadcast_three_input",
       "priority": 20,
+      "when": ["variadicInputCount == 3", "ranks.a <= ranks.y", "ranks.b <= ranks.y", "ranks.c <= ranks.y", "broadcastOutputOk"],
       "passes": [
         {
           "id": "main",
           "name": "Mean",
+          "shader": "summean-broadcast.wgsl.jinja",
+          "derive": {
+            "aShape": "shapes.a",
+            "bShape": "shapes.b",
+            "cShape": "shapes.c",
+            "yShape": "shapes.y",
+            "aRank": "ranks.a",
+            "bRank": "ranks.b",
+            "cRank": "ranks.c",
+            "yRank": "ranks.y",
+            "hasC": "true",
+            "op": "\"mean\""
           },
+          "bindings": [{ "arg": "a" }, { "arg": "b" }, { "arg": "c" }, "y", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
       "id": "broadcast_four_input_vec4",
       "priority": 26,
       "when": ["variadicInputCount == 4", "broadcastVec4OutputOk", "aBroadcastVec4Ok", "bBroadcastVec4Ok", "cBroadcastVec4Ok", "dBroadcastVec4Ok"],
+      "derive": {
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "aElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if ranks.a > 0 and dim(shapes.a, ranks.a - 1) != 1 else dtypes.T",
+        "bElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if ranks.b > 0 and dim(shapes.b, ranks.b - 1) != 1 else dtypes.T",
+        "cElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if ranks.c > 0 and dim(shapes.c, ranks.c - 1) != 1 else dtypes.T",
+        "dElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if ranks.d > 0 and dim(shapes.d, ranks.d - 1) != 1 else dtypes.T"
       },
       "passes": [
         {
           "id": "main",
           "name": "Mean.Broadcast4Vec4",
+          "shader": "summean-broadcast.wgsl.jinja",
+          "derive": {
+            "aShape": "shapes.a",
+            "bShape": "shapes.b",
+            "cShape": "shapes.c",
+            "dShape": "shapes.d",
+            "yShape": "shapes.y",
+            "aRank": "ranks.a",
+            "bRank": "ranks.b",
+            "cRank": "ranks.c",
+            "dRank": "ranks.d",
+            "yRank": "ranks.y",
+            "aVector": "ranks.a > 0 and dim(shapes.a, ranks.a - 1) != 1",
+            "bVector": "ranks.b > 0 and dim(shapes.b, ranks.b - 1) != 1",
+            "cVector": "ranks.c > 0 and dim(shapes.c, ranks.c - 1) != 1",
+            "dVector": "ranks.d > 0 and dim(shapes.d, ranks.d - 1) != 1",
+            "hasC": true,
+            "hasD": true,
+            "vectorizedSpec": true,
+            "op": "\"mean\""
           },
+          "bindings": ["a_3", "b_3", "c_2", { "arg": "d", "elementType": "$dElement" }, "y_2", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "broadcast_four_input",
       "priority": 22,
+      "when": ["variadicInputCount == 4", "ranks.a <= ranks.y", "ranks.b <= ranks.y", "ranks.c <= ranks.y", "ranks.d <= ranks.y", "broadcastOutputOk"],
       "passes": [
         {
           "id": "main",
           "name": "Mean",
+          "shader": "summean-broadcast.wgsl.jinja",
+          "derive": {
+            "aShape": "shapes.a",
+            "bShape": "shapes.b",
+            "cShape": "shapes.c",
+            "yShape": "shapes.y",
+            "aRank": "ranks.a",
+            "bRank": "ranks.b",
+            "cRank": "ranks.c",
+            "yRank": "ranks.y",
+            "hasC": "true",
+            "op": "\"mean\"",
+            "dShape": "shapes.d",
+            "dRank": "ranks.d",
+            "hasD": "true"
           },
+          "bindings": [{ "arg": "a" }, { "arg": "b" }, { "arg": "c" }, { "arg": "d" }, "y", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "broadcast_five_input",
       "priority": 23,
+      "when": ["variadicInputCount == 5", "ranks.a <= ranks.y", "ranks.b <= ranks.y", "ranks.c <= ranks.y", "ranks.d <= ranks.y", "ranks.e <= ranks.y", "broadcastOutputOk"],
       "passes": [
         {
           "id": "main",
           "name": "Mean5",
+          "shader": "summean-broadcast.wgsl.jinja",
+          "derive": {
+            "aShape": "shapes.a",
+            "bShape": "shapes.b",
+            "cShape": "shapes.c",
+            "yShape": "shapes.y",
+            "aRank": "ranks.a",
+            "bRank": "ranks.b",
+            "cRank": "ranks.c",
+            "yRank": "ranks.y",
+            "hasC": "true",
+            "op": "\"mean\"",
+            "dShape": "shapes.d",
+            "dRank": "ranks.d",
+            "hasD": "true",
+            "hasE": "true",
+            "extraInputs": "[\"c\", \"d\", \"e\"]",
+            "eShape": "shapes.e",
+            "eRank": "ranks.e"
           },
+          "bindings": [{ "arg": "a" }, { "arg": "b" }, { "arg": "c" }, { "arg": "d" }, "e", "y", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,20 +1,35 @@
 {
   "name": "ai.onnx.Mean",
-  "id": "_ai_onnx_mean_webgpu_6e7e02e",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "PZuRYEh9H7m59RtG7mKg5vKX61mxb+faVyXw45NmJkg=",
-      "datamove-elementwise-copy.wgsl.jinja": "J5yC2bAddPiP+odXLgVGS3TJ9jeNsfRTvedKrj/fhZg=",
-      "manifest.json": "VktTxdT7xm0Gpy3RAjtJt7o1jS2zDQaXet8yeeO27bY=",
-      "summean-broadcast.wgsl.jinja": "JV9W3CZ+lxla2C9yTyminEV2qBzBys8/q1rwdBcSTkA=",
-      "summean-vec4.wgsl.jinja": "mZSVsrneREFgo5NJjDSgjU5zmMhMD91OA39Z2i+5xlM=",
-      "test.json": "mpFM+ziMgNaNEpbiAD6n5ZN5MWRx1JiuISOPJsgfeLw="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Mean" }
 }

 {
   "name": "ai.onnx.Mean",
+  "id": "_ai_onnx_mean_webgpu_1ff9069",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "OCGno3NWR0wjZ5xDoYp8Q1ELrQ2KCph+dNK9rQYSZzk=",
+      "datamove-elementwise-copy.wgsl.jinja": "Q1WuZCqcDf9b6cbT4rLilpeWj6327t/bSFmKr6IVBoQ=",
+      "manifest.json": "07k7OHtxzThVNzb5aKOjXpdsovnIZCCdqNqyHtzJDl8=",
+      "summean-broadcast.wgsl.jinja": "WPpxIr6ztCcEDvyREMrLJVVWnu3GKvhH8tzIteHMKGU=",
+      "summean-vec4.wgsl.jinja": "oNlvKJctx7Gn2D6aHpTcVsf+ZyoYEadtFgyVQZV+8So=",
+      "test.json": "/isBt5ibLwpR8jqM0J4B1cqABoi9FeFoLOE7YxZwRgg="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "single_input_identity": ["datamove-elementwise-copy.wgsl.jinja"],
+      "same_shape_vec4_two_input": ["summean-vec4.wgsl.jinja"],
+      "same_shape_vec4_three_input": ["summean-vec4.wgsl.jinja"],
+      "same_shape_vec4_four_input": ["summean-vec4.wgsl.jinja"],
+      "same_shape_vec4_five_input": ["summean-vec4.wgsl.jinja"],
+      "broadcast_two_input": ["summean-broadcast.wgsl.jinja"],
+      "broadcast_three_input_vec4": ["summean-broadcast.wgsl.jinja"],
+      "broadcast_three_input": ["summean-broadcast.wgsl.jinja"],
+      "broadcast_four_input_vec4": ["summean-broadcast.wgsl.jinja"],
+      "broadcast_four_input": ["summean-broadcast.wgsl.jinja"],
+      "broadcast_five_input": ["summean-broadcast.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/summean-broadcast.wgsl.jinja CHANGED Viewed

@@ -1,8 +1,5 @@
-{% set allInputs = ["a", "b"] + (source.extraInputs if source.extraInputs is defined else (["c"] if source.hasC else []) + (["d"] if source.hasD else [])) %}
-{% set extraInputs = source.extraInputs if source.extraInputs is defined else (["c"] if source.hasC else []) + (["d"] if source.hasD else []) %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% macro offset_fn(fn_name, opShape, opRank, op_same, op_numel, outShape, outRank, out_numel) %}
 fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif %}) -> u32 {
@@ -68,34 +65,34 @@ fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif
-{{ broadcast_offset_fn("a_offset", source.aShape, source.aRank, source.yShape, source.yRank) }}
-{{ broadcast_offset_fn("b_offset", source.bShape, source.bRank, source.yShape, source.yRank) }}
 {% for n in extraInputs %}
-{{ broadcast_offset_fn(n ~ "_offset", source[n ~ "Shape"], source[n ~ "Rank"], source.yShape, source.yRank) }}
 {% endfor %}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
-fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit.
-  let i = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
   if (i >= params.count) {
     return;
   }
-{% if source.vectorized %}
   // The output's innermost dimension is four-aligned. Each input either keeps
   // that dimension (one aligned vec4 load) or broadcasts it (one scalar splat).
   // Offset folding is therefore paid once per four outputs without changing
   // multidirectional broadcast semantics on any outer dimension.
   let base = i * 4u;
 {% for n in allInputs %}
-{% if source[n ~ "Vector"] %}
-  let {{ n }}v = vec4<f32>({{ n }}[{{ broadcast_offset_call(n ~ "_offset", source[n ~ "Shape"], source.yShape, "base") }} / 4u]);
 {% else %}
-  let {{ n }}v = vec4<f32>(f32({{ n }}[{{ broadcast_offset_call(n ~ "_offset", source[n ~ "Shape"], source.yShape, "base") }}]));
 {% endif %}
 {% endfor %}
   let total = {% for n in allInputs %}{{ n }}v{% if not loop.last %} + {% endif %}{% endfor %};
@@ -105,7 +102,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
   y[i] = {{ vectorScalar }}(total / {{ allInputs | length }}.0);
 {% endif %}
 {% else %}
-  let total = {% for n in allInputs %}f32({{ n }}[{{ broadcast_offset_call(n ~ "_offset", source[n ~ "Shape"], source.yShape, "i") }}]){% if not loop.last %} + {% endif %}{% endfor %};
 {% if allInputs | length == 2 %}
   y[i] = {{ scalar }}(0.5 * total);
 {% else %}

+{% set allInputs = ["a", "b"] + (extraInputs if extraInputs is defined else (["c"] if hasC else []) + (["d"] if hasD else [])) %}
+{% set extraInputs = extraInputs if extraInputs is defined else (["c"] if hasC else []) + (["d"] if hasD else []) %}
 {{ env.wgsl.resourceDeclarations }}
 {% macro offset_fn(fn_name, opShape, opRank, op_same, op_numel, outShape, outRank, out_numel) %}
 fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif %}) -> u32 {
+{{ broadcast_offset_fn("a_offset", aShape, aRank, yShape, yRank) }}
+{{ broadcast_offset_fn("b_offset", bShape, bRank, yShape, yRank) }}
 {% for n in extraInputs %}
+{{ broadcast_offset_fn(n ~ "_offset", derive[n ~ "Shape"], derive[n ~ "Rank"], yShape, yRank) }}
 {% endfor %}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits past the
+  // per-axis dispatch fold width.
+  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if (i >= params.count) {
     return;
   }
+{% if vectorizedSpec %}
   // The output's innermost dimension is four-aligned. Each input either keeps
   // that dimension (one aligned vec4 load) or broadcasts it (one scalar splat).
   // Offset folding is therefore paid once per four outputs without changing
   // multidirectional broadcast semantics on any outer dimension.
   let base = i * 4u;
 {% for n in allInputs %}
+{% if derive[n ~ "Vector"] %}
+  let {{ n }}v = vec4<f32>({{ n }}[{{ broadcast_offset_call(n ~ "_offset", derive[n ~ "Shape"], yShape, "base") }} / 4u]);
 {% else %}
+  let {{ n }}v = vec4<f32>(f32({{ n }}[{{ broadcast_offset_call(n ~ "_offset", derive[n ~ "Shape"], yShape, "base") }}]));
 {% endif %}
 {% endfor %}
   let total = {% for n in allInputs %}{{ n }}v{% if not loop.last %} + {% endif %}{% endfor %};
   y[i] = {{ vectorScalar }}(total / {{ allInputs | length }}.0);
 {% endif %}
 {% else %}
+  let total = {% for n in allInputs %}f32({{ n }}[{{ broadcast_offset_call(n ~ "_offset", derive[n ~ "Shape"], yShape, "i") }}]){% if not loop.last %} + {% endif %}{% endfor %};
 {% if allInputs | length == 2 %}
   y[i] = {{ scalar }}(0.5 * total);
 {% else %}

build/webgpu/summean-vec4.wgsl.jinja CHANGED Viewed

@@ -1,15 +1,12 @@
-{% set extraInputs = source.extraInputs if source.extraInputs is defined else (["c"] if source.hasC else []) + (["d"] if source.hasD else []) %}
 {% set arity = 2 + extraInputs | length %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
-fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit.
-  let i = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
   if (i >= params.count) {
     return;
   }

+{% set extraInputs = extraInputs if extraInputs is defined else (["c"] if hasC else []) + (["d"] if hasD else []) %}
 {% set arity = 2 + extraInputs | length %}
 {{ env.wgsl.resourceDeclarations }}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits past the
+  // per-axis dispatch fold width.
+  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if (i >= params.count) {
     return;
   }

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.Mean",
   "cases": [
     {
       "name": "max_arity_float16_positions",
@@ -71,7 +70,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc",
         "test": "MathOpTest.Mean_6",
-        "notes": "Vec4 companion: the mean of equal finite subnormal values should remain subnormal in the vectorized path."
       },
       "inputs": {
         "a": {
@@ -151,7 +150,7 @@
     {
       "name": "float16_vec4_three_input",
       "provenance": {
-        "notes": "Route lock for the f16 three-input vec4 path. the summean-vec4 kernel has a separate f16 branch that widens each operand to f32, accumulates, then narrows on store; until now only two-input f16 renders existed, so on that branch neither the extra-operand loop that folds C in nor the non-arity-2 divisor (total / 3.0, as opposed to the 0.5 * total shortcut the two-input case takes) was ever emitted. Both numerator and denominator are therefore checked here: with C dropped, or with the 0.5 shortcut applied to three operands, every lane moves. Every operand, every partial sum and every quotient is exactly representable in float16 (each total is 3x a dyadic value), so the expected output is the exact (a + b + c) / 3 and needs no tolerance."
       },
       "inputs": {
         "a": {
@@ -233,7 +232,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc",
         "test": "MathOpTest.Mean_6",
-        "notes": "Extends ORT's same-shape Mean coverage to a valid four-input ONNX variadic node."
       },
       "inputs": {
         "a": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "constant", "value": 1.0 } },
@@ -248,7 +247,7 @@
     {
       "name": "four_input_same_shape_vec4_reference_generated",
       "provenance": {
-        "notes": "Reference-generated four-operand coverage. Mean was doubly exposed: reducing only A..C also divided by three, so both the numerator and the denominator were wrong. Every element moves if either is, which makes this a total detector rather than a partial one."
       },
       "inputs": {
         "a": {
@@ -279,7 +278,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc",
         "test": "MathOpTest.Mean_6",
-        "notes": "Extends ORT's variadic Mean coverage with a fourth input that exposes NaN propagation and inf-minus-inf cancellation."
       },
       "inputs": {
         "a": { "dtype": "float32", "shape": [4], "data": { "kind": "values", "values": [1.0, "Infinity", 1.0, 8.0] } },
@@ -317,7 +316,7 @@
     {
       "name": "three_input_broadcast_vec4_mean",
       "provenance": {
-        "notes": "Reaches broadcast_three_input_vec4. The output's innermost dimension is four-aligned, `a` keeps it (one aligned vec4 load) and `b` and `c` broadcast it (scalar splats), so the variant's three per-input element types are not all the same -- which is what the vec4 broadcast binding set exists for. Its four-input sibling was already reachable from four_input_scalar_d_broadcast_mean; this one had nothing."
       },
       "inputs": {
         "a": {

 {
   "cases": [
     {
       "name": "max_arity_float16_positions",
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc",
         "test": "MathOpTest.Mean_6",
+        "notes": "On the vec4 path, the mean of equal finite subnormal values must remain subnormal."
       },
       "inputs": {
         "a": {
     {
       "name": "float16_vec4_three_input",
       "provenance": {
+        "notes": "Three float16 vec4 inputs exercise widening each operand to float32, folding the third input, dividing by three, and narrowing on store. Dyadic inputs make every operand, partial sum, and quotient exactly representable in float16, so the expected result is exactly (a+b+c)/3."
       },
       "inputs": {
         "a": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc",
         "test": "MathOpTest.Mean_6",
+        "notes": "Four same-shaped inputs exercise the maximum supported arity of this variadic Mean package."
       },
       "inputs": {
         "a": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "constant", "value": 1.0 } },
     {
       "name": "four_input_same_shape_vec4_reference_generated",
       "provenance": {
+        "notes": "Four distinct operands make both the numerator and divisor observable: omitting the fourth input or dividing by three changes every output element."
       },
       "inputs": {
         "a": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc",
         "test": "MathOpTest.Mean_6",
+        "notes": "A fourth input exposes NaN propagation and positive-Infinity plus negative-Infinity cancellation."
       },
       "inputs": {
         "a": { "dtype": "float32", "shape": [4], "data": { "kind": "values", "values": [1.0, "Infinity", 1.0, 8.0] } },
     {
       "name": "three_input_broadcast_vec4_mean",
       "provenance": {
+        "notes": "The four-aligned innermost output dimension permits a vec4 load from a while b and c broadcast as scalar splats, exercising three-input vectorized broadcasting with mixed binding element types."
       },
       "inputs": {
         "a": {