| { |
| "domain": "ai.onnx", |
| "name": "BatchNormalization", |
| "sinceVersion": 15, |
| "description": "Applies inference-mode batch normalization: `Y = (X - input_mean) / sqrt(input_var + epsilon) * scale + B`. This package supports `training_mode=0`, rank-2-or-higher inputs, and a common float16 or float32 dtype for every tensor. ONNX training mode is intentionally not implemented because this inference-only release does not expose its required running-mean and running-variance outputs.", |
| "inputs": [ |
| { |
| "role": "X", |
| "dtype": "T", |
| "description": "Input data tensor with shape `(N, C, D1, ..., Dn)`, normalized independently per channel using the supplied estimated statistics." |
| }, |
| { "role": "scale", "dtype": "T", "rank": 1, "description": "Per-channel scale tensor with shape `(C)`." }, |
| { "role": "B", "dtype": "T", "rank": 1, "description": "Per-channel bias tensor with shape `(C)`." }, |
| { |
| "role": "input_mean", |
| "dtype": "T", |
| "rank": 1, |
| "description": "Precomputed estimated mean tensor with shape `(C)` used for inference." |
| }, |
| { |
| "role": "input_var", |
| "dtype": "T", |
| "rank": 1, |
| "description": "Precomputed estimated variance tensor with shape `(C)` used for inference." |
| } |
| ], |
| "outputs": [ |
| { |
| "role": "Y", |
| "dtype": "T", |
| "rank": "ranks.X", |
| "description": "Batch-normalized output tensor with the same shape as `X`.", |
| "shape": "shapes.X" |
| } |
| ], |
| "attributes": { "epsilon": 0.00001, "momentum": 0.9, "training_mode": 0 }, |
| "attributeDescriptions": { |
| "epsilon": "Small value added to the variance before taking the square root to avoid division by zero.", |
| "momentum": "Standard ONNX running-statistics momentum. This inference-only package accepts the default `0.9`; non-default values are reserved for the unsupported training-state update.", |
| "training_mode": "Execution mode. This inference-only package supports the default value 0; value 1 is rejected because the ONNX training outputs are not exposed." |
| }, |
| "attributeConstraints": { |
| "momentum": { "values": [0.9], "comparison": "float32" }, |
| "training_mode": { "values": [0] } |
| }, |
| "typeConstraints": { "T": ["float32", "float16"] }, |
| "args": { |
| "x": { "kind": "tensor", "semantic": "X", "role": "input" }, |
| "scale": { "kind": "tensor", "semantic": "scale", "role": "input" }, |
| "b": { "kind": "tensor", "semantic": "B", "role": "input" }, |
| "inputMean": { "kind": "tensor", "semantic": "input_mean", "role": "input" }, |
| "inputVar": { "kind": "tensor", "semantic": "input_var", "role": "input" }, |
| "y": { "kind": "tensor", "semantic": "Y", "role": "output" } |
| }, |
| "tunables": { "WORKGROUP_SIZE": 256 }, |
| "derive": { |
| "normalizationParamsOk": "ranks.scale == 1 and ranks.B == 1 and ranks.input_mean == 1 and ranks.input_var == 1 and dim(shapes.scale, 0) == dim(shapes.X, 1) and dim(shapes.B, 0) == dim(shapes.X, 1) and dim(shapes.input_mean, 0) == dim(shapes.X, 1) and dim(shapes.input_var, 0) == dim(shapes.X, 1)", |
| "inferenceContractOk": "f16Ok(dtypes.T) and ranks.X >= 2 and ranks.Y == ranks.X and sameShape(shapes.Y, shapes.X) and normalizationParamsOk" |
| }, |
| "bindingSets": { |
| "ncInferenceVec4": [ |
| { |
| "name": "x", |
| "arg": "x", |
| "semantic": "X", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "vec4<f32>" |
| }, |
| { |
| "name": "scale", |
| "arg": "scale", |
| "semantic": "scale", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "vec4<f32>" |
| }, |
| { |
| "name": "bias", |
| "arg": "b", |
| "semantic": "B", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "vec4<f32>" |
| }, |
| { |
| "name": "input_mean", |
| "arg": "inputMean", |
| "semantic": "input_mean", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "vec4<f32>" |
| }, |
| { |
| "name": "input_var", |
| "arg": "inputVar", |
| "semantic": "input_var", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "vec4<f32>" |
| }, |
| { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "vec4<f32>" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "count4", "type": "u32", "value": "numel(shapes.Y) / 4" }, |
| { "name": "channels4", "type": "u32", "value": "dim(shapes.X, 1) / 4" }, |
| { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" } |
| ] |
| } |
| } |
| ], |
| "spatialInferenceVec4": [ |
| { |
| "name": "x", |
| "arg": "x", |
| "semantic": "X", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "vec4<f32>" |
| }, |
| { |
| "name": "scale", |
| "arg": "scale", |
| "semantic": "scale", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$T" |
| }, |
| { "name": "bias", "arg": "b", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, |
| { |
| "name": "input_mean", |
| "arg": "inputMean", |
| "semantic": "input_mean", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$T" |
| }, |
| { |
| "name": "input_var", |
| "arg": "inputVar", |
| "semantic": "input_var", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$T" |
| }, |
| { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "vec4<f32>" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "count4", "type": "u32", "value": "numel(shapes.Y) / 4" }, |
| { "name": "spatial4", "type": "u32", "value": "inner(shapes.X, 1) / 4" }, |
| { "name": "channels", "type": "u32", "value": "dim(shapes.X, 1)" }, |
| { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" } |
| ] |
| } |
| } |
| ], |
| "inferenceScalar": [ |
| { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, |
| { |
| "name": "scale", |
| "arg": "scale", |
| "semantic": "scale", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$T" |
| }, |
| { "name": "bias", "arg": "b", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, |
| { |
| "name": "input_mean", |
| "arg": "inputMean", |
| "semantic": "input_mean", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$T" |
| }, |
| { |
| "name": "input_var", |
| "arg": "inputVar", |
| "semantic": "input_var", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$T" |
| }, |
| { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$T" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "channels", "type": "u32", "value": "dim(shapes.X, 1)" }, |
| { "name": "height", "type": "u32", "value": "1 if ranks.X == 2 else dim(shapes.X, 2)" }, |
| { "name": "width", "type": "u32", "value": "1 if ranks.X == 2 else inner(shapes.X, 2)" }, |
| { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" }, |
| { "name": "count", "type": "u32", "value": "numel(shapes.Y)" } |
| ] |
| } |
| } |
| ] |
| }, |
| "variants": [ |
| { |
| "id": "inference_scalar", |
| "when": ["inferenceContractOk"], |
| "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "BatchNormalization.InferenceScalar", |
| "shader": "batch-normalization-nchw.wgsl.jinja", |
| "bindings": "inferenceScalar", |
| "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": "tunables.WORKGROUP_SIZE" } |
| } |
| ] |
| }, |
| { |
| "id": "nc_inference_vec4", |
| "priority": 110, |
| "when": ["inferenceContractOk", "dtypes.T == \"f32\"", "ranks.X == 2", "dim(shapes.X, 1) % 4 == 0"], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "BatchNormalization.NcInferenceVec4", |
| "shader": "batch-normalization-nc-vec4.wgsl.jinja", |
| "bindings": "ncInferenceVec4", |
| "dispatch": { "threads": "numel(shapes.Y) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" } |
| } |
| ] |
| }, |
| { |
| "id": "nchw_inference_vec4", |
| "priority": 100, |
| "when": ["inferenceContractOk", "dtypes.T == \"f32\"", "ranks.X >= 3", "inner(shapes.X, 1) % 4 == 0"], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "BatchNormalization.InferenceVec4", |
| "shader": "batch-normalization-nchw-vec4.wgsl.jinja", |
| "bindings": "spatialInferenceVec4", |
| "dispatch": { "threads": "numel(shapes.Y) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" } |
| } |
| ] |
| } |
| ] |
| } |
|
|
