library_name: kernels
license: apache-2.0
tags:
- kernel
- webgpu
- wgsl
ai.onnx.ArgMax
ai.onnx · standard ONNX operator · ONNX opset ≥ 13
Description
Returns the index of the maximum value along an axis, choosing the first equal value unless select_last_index is enabled.
See the ONNX ArgMax spec for the reference semantics.
Inputs
| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
|---|---|---|---|---|---|---|
data |
x |
T |
— | — | Values whose maximum index is selected along axis. |
required |
Outputs
| Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
|---|---|---|---|---|---|---|---|
reduced |
y |
I |
uint32 |
derived | derived; see description | Logical int64 indices of the maximum values along the reduced axis; WebGPU stores these bounded indices as uint32. | required |
Attributes
Default values (overridable per request):
| Attribute | Default | Description |
|---|---|---|
axis |
0 |
Axis to reduce; negative values count from the back. |
keepdims |
1 |
Retain the reduced dimension with length one when non-zero. |
select_last_index |
0 |
Choose the last equal maximum instead of the first when non-zero. |
Type constraints
| Variable | Allowed dtypes |
|---|---|
T |
float32, float16, int32, uint32, int16, int8, uint8 |
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 casesbench.json— benchmark + tuning casesreduce-arg-axis-split-combine.wgsl.jinjareduce-arg-axis-split-reduce.wgsl.jinjareduce-arg-axis-split-tiled.wgsl.jinjareduce-arg-axis-tiled.wgsl.jinjareduce-arg-axis.wgsl.jinjareduce-arg-row-split.wgsl.jinjareduce-arg-row-subgroup.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.ArgMax", { version: 1 });
const { y } = await kernel({ x: { data: xData, shape: [2, 2] } });