ai.onnx.RNN

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

Description

Computes a one-layer simple RNN, updating the hidden state each time step as Ht = f(Xt*(Wi^T) + Ht-1*(Ri^T) + Wbi + Rbi) where f defaults to Tanh. Supports forward, reverse, and bidirectional traversal; optional clip bounds pre-activation inputs to [-clip, +clip].

See the ONNX RNN spec for the reference semantics.

Inputs

Name	Bind key	Logical dtype	Rank	Shape	Description	Presence
`X`	`x`	`T`	`3`	—	Input sequences with shape `[seq_length, batch_size, input_size]` when `layout=0`, or `[batch_size, seq_length, input_size]` when `layout=1`.	required
`W`	`w`	`T`	`3`	—	Weight tensor for the input gate, shape `[num_directions, hidden_size, input_size]`.	required
`R`	`r`	`T`	`3`	—	Recurrence weight tensor, shape `[num_directions, hidden_size, hidden_size]`.	required
`B`	`b`	`T`	`2`	—	Bias tensor concatenating input and recurrence biases, shape `[num_directions, 2*hidden_size]`. ONNX defines an omitted bias as zero, but this kernel ABI requires an explicit tensor; callers representing omission must bind a zero-filled tensor.	required
`sequence_lens`	`sequence_lens`	`int32`	`1`	—	Per-batch sequence lengths of shape `[batch_size]`; all sequences are treated as full length if absent.	optional
`initial_h`	`initial_h`	`T`	`3`	—	Initial hidden state with shape `[num_directions, batch_size, hidden_size]` when `layout=0`, or `[batch_size, num_directions, hidden_size]` when `layout=1`; assumed zero if absent.	optional

Outputs

Name	Bind key	Logical dtype	Rank	Shape	Description	Presence
`Y`	`y`	`T`	`4`	derived; see description	All intermediate hidden outputs, with shape `[seq_length, num_directions, batch_size, hidden_size]` when `layout=0`, or `[batch_size, seq_length, num_directions, hidden_size]` when `layout=1`.	required
`Y_h`	`y_h`	`T`	`3`	derived; see description	Final hidden state with shape `[num_directions, batch_size, hidden_size]` when `layout=0`, or `[batch_size, num_directions, hidden_size]` when `layout=1`.	required

Attributes

Attributes and default values (overridable per request):

Attribute	Default	Description
`layout`	`0`	Shape format for `X`, `initial_h`, `Y`, and `Y_h`: `0` puts `seq_length` first; `1` puts `batch_size` first.
`direction`	`"forward"`	Traversal direction: `"forward"`, `"reverse"`, or `"bidirectional"`.
`activations`	`["Tanh","Tanh"]`	Activation function for each direction. ONNX defaults both directions to `Tanh`; `forward` and `reverse` use the first entry, while `bidirectional` execution uses both entries.
`activation_alpha`	—	Optional alpha parameters for activation functions that use them, consumed in activation order; omitted entries use each activation's standard default.
`activation_beta`	—	Optional beta parameters for activation functions that use them, consumed in activation order; omitted entries use each activation's standard default.
`clip`	—	Optional non-negative threshold applied to pre-activation values as `[-clip, +clip]`; omission disables clipping, while an explicit `0` clamps them to zero.
`hidden_size`	—	Optional number of neurons in the hidden layer; when omitted, it is inferred from the W and R tensor shapes.

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
datamove-transpose-2d-tiled.wgsl.jinja
dense-tiled-matmul.wgsl.jinja
rnn-general.wgsl.jinja
rnn-input-proj.wgsl.jinja
rnn-parallel-global.wgsl.jinja
rnn-parallel.wgsl.jinja
rnn-recur-step.wgsl.jinja
rnn-recur-subgroup-step.wgsl.jinja
rnn-sequence-mask.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.RNN", { version: 1 });
const { y, y_h } = await kernel({
  x: { data: xData, shape: [1, 1, 1] },
  w: { data: wData, shape: [1, 1, 1] },
  r: { data: rData, shape: [1, 1, 1] },
  b: { data: bData, shape: [1, 2] },
});

Downloads last month: -

kernel

webgpu

wgsl

apache-2.0

WebGPU

Requires WebGPU support. See the compatibility table.