ai.onnx.TopK

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

Description

Retrieves the top-k largest or smallest elements along the selected axis, returning values and stable lower-index tie-breaking indices.

See the ONNX TopK spec for the reference semantics.

Inputs

Name	Bind key	Logical dtype	Rank	Shape	Description	Presence
`X`	`x`	`T`	—	—	Values from which the top `k` entries are selected along `axis`.	required

Outputs

Name	Bind key	Logical dtype	WebGPU storage	Rank	Shape	Description	Presence
`Values`	`values`	`T`	runtime-selected; narrow integers and bool use 32-bit slots	same as `X`	derived; see description	Selected values; the reduced axis has length `k`.	required
`Indices`	`indices`	`I`	`uint32`	same as `X`	derived; see description	Logical int64 indices of the selected values along the reduced axis; WebGPU stores these bounded indices as uint32.	required

Runtime arguments

Name	Kind	Semantic	Description	Presence
`k`	`u32`	`kernel.k`	Number of values to select along the configured axis.	required

Attributes

Default values (overridable per request):

Attribute	Default	Description
`axis`	`-1`	Axis to reduce; negative values count from the back.
`largest`	`1`	Select largest values when 1, smallest values when 0.
`sorted`	`1`	Sort selected values when 1. A sorted result is also valid when output order is unspecified (`sorted=0`).

Type constraints

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

Device requirements

Some implementation variants require subgroups. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.

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
topk-axis.wgsl.jinja
topk-large-block.wgsl.jinja
topk-noop.wgsl.jinja
topk-portable-rows-smallk.wgsl.jinja
topk-small-rows-batched.wgsl.jinja
topk-strided-smallk.wgsl.jinja
topk-subgroup-rows.wgsl.jinja
topk-top1-last-axis.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.TopK", { version: 1 });
const { values, indices } = await kernel({ x: { data: xData, shape: [1, 3] }, k: 1 });

Downloads last month: -

kernel

webgpu

wgsl

apache-2.0

WebGPU

Requires WebGPU support. See the compatibility table.