{ "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" }, { "name": "scale", "arg": "scale", "semantic": "scale", "buffer": { "type": "read-only-storage" }, "elementType": "vec4" }, { "name": "bias", "arg": "b", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "vec4" }, { "name": "input_mean", "arg": "inputMean", "semantic": "input_mean", "buffer": { "type": "read-only-storage" }, "elementType": "vec4" }, { "name": "input_var", "arg": "inputVar", "semantic": "input_var", "buffer": { "type": "read-only-storage" }, "elementType": "vec4" }, { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "vec4" }, { "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" }, { "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" }, { "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" } } ] } ] }