ai.onnx.Scan

ai.onnx · internal tensor lowering (non-standard) · reviewed against ONNX opset 25

Description

Support status: the standard variadic ONNX Scan control-flow operator is not implemented because standalone kernel packages cannot carry or execute its body graph. This internal lowering computes a forward or reverse additive prefix scan for one rank-1 state and one rank-2 input and must not be treated as ONNX Scan.

See the standard ONNX Scan spec for the contract this internal lowering does not implement.

Inputs

Name	Bind key	Logical dtype	Rank	Shape	Description	Presence
`initial_state`	`initial_state`	`T`	`1`	—	Initial rank-1 additive state of shape `[dim]`.	required
`scan_input`	`scan_input`	`T`	`2`	—	Rank-2 input of shape `[steps, dim]`. Each row is added to the running state in forward or reverse traversal order.	required

Outputs

Name	Bind key	Logical dtype	Rank	Shape	Description	Presence
`final_state`	`final_state`	`T`	`1`	same as `initial_state`	Final rank-1 state of shape `[dim]` after all input rows have been accumulated.	required
`scan_output`	`scan_output`	`T`	`2`	same as `scan_input`	Inclusive additive prefix results with the same `[steps, dim]` shape as `scan_input`. Reverse traversal still writes each result at its corresponding input row.	required

Attributes

Default values (overridable per request):

Attribute	Default	Description
`reverse`	`0`	When non-zero, the scan input sequence is traversed in reverse order (equivalent to `scan_input_directions=1`); default `0` scans forward.

Type constraints

Variable	Allowed dtypes
`T`	`float32`

Files

metadata.json — kernel metadata (id, digests, provenance)
manifest.json — the op contract (source of truth)
test.json — correctness cases
bench.json — benchmark + tuning cases
scan-coop-channel.wgsl.jinja
scan-multichunk-apply.wgsl.jinja
scan-multichunk-carries.wgsl.jinja
scan-multichunk-local.wgsl.jinja
scan-prefix-sum.wgsl.jinja

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.

import { getKernel } from "@huggingface/kernels";

const kernel = await getKernel("webgpu-kernels/ai.onnx.Scan", { version: 1 });
const { final_state, scan_output } = await kernel({
  initial_state: { data: initial_stateData, shape: [1] },
  scan_input: { data: scan_inputData, shape: [3, 1] },
});

Downloads last month: -

kernel

webgpu

wgsl

apache-2.0

WebGPU

Requires WebGPU support. See the compatibility table.