ai.onnx.Split / README.md

Xenova HF Staff

sync 2e7068faf55e

854303d verified 6 days ago

3.94 kB

library_name: kernels
license: apache-2.0
tags:
  - kernel
  - webgpu
  - wgsl

ai.onnx.Split

ai.onnx · standard ONNX operator · ONNX opset ≥ 18

Description

Splits a tensor into a list of tensors along the specified axis. The sizes of each output slice are given by the optional split input; if omitted, the tensor is divided into equal parts (the last chunk may be smaller if the axis dimension is not evenly divisible).

See the ONNX Split spec for the reference semantics.

Inputs

Name	Bind key	Logical dtype	WebGPU storage	Rank	Shape	Description	Presence
`input`	`input`	`T`	runtime-selected; narrow integers and bool use 32-bit slots	—	—	The tensor to split.	required
`split`	`split`	`S`	`uint32`	`1`	—	Optional logical int64 1-D tensor specifying the size of each output along the split axis; values must be non-negative, sum to the axis dimension, and use uint32 WebGPU storage.	optional

Outputs

Name	Bind key	Logical dtype	Rank	Shape	Description	Presence
`Y0`	`y0`	`T`	same as `input`	—	First output slice after splitting.	required
`Y1`	`y1`	`T`	same as `input`	—	Second output slice after splitting (optional).	optional
`Y2`	`y2`	`T`	same as `input`	—	Third output slice after splitting.	optional
`Y3`	`y3`	`T`	same as `input`	—	Fourth output slice after splitting.	optional
`Y4`	`y4`	`T`	same as `input`	—	Fifth output slice after splitting.	optional

Attributes

Attributes and default values (overridable per request):

Attribute	Default	Description
`axis`	`0`	The axis along which to split. Negative values count from the end; accepted range is `[-rank, rank-1]`.
`num_outputs`	—	Optional number of outputs when the `split` input is omitted. The final output may be smaller when the axis dimension is not evenly divisible.

Type constraints

Variable	Allowed dtypes
`T`	`float32`, `float16`, `uint32`, `int32`, `int16`, `uint8`, `int8`, `bool`
`S`	`int64`

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
datamove-flat-copy.wgsl.jinja
datamove-split-block.wgsl.jinja
split-n.wgsl.jinja

Use with `@huggingface/kernels`

The loader automatically allocates outputs whose metadata it can derive from the manifest contract and this call.

The explicit outputs entries provide shape and logical dtype metadata for the results listed below:

y0
y1
y2

Each entry either requests an optional result or supplies metadata that cannot be inferred from the inputs.

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.Split", { version: 1 });
// Explicit destinations request optional results or supply metadata that cannot be inferred.
const { y0, y1, y2 } = await kernel({ input: { data: inputData, shape: [6] } }, {
  outputs: {
    y0: { shape: [2], dtype: "float32" },
    y1: { shape: [2], dtype: "float32" },
    y2: { shape: [2], dtype: "float32" },
  },
});