include_nn_kernels.h

namespace onnx_light

Alias that makes onnx-light headers compatible with code that references ONNX_LIGHT_NAMESPACE (the macro used in the standard onnx package).

Set to ONNX_LIGHT_NAMESPACE so both names resolve to the same namespace.

Symbol-visibility attribute for the public onnx-light C++ API.

Defined as empty because onnx-light does not require explicit __declspec(dllexport) or __attribute__((visibility("default"))) annotations — visibility is controlled at the shared-library level. The macro is provided so that vendored ONNX headers that decorate their declarations with ONNX_API compile without modification.

Namespace alias so that ONNX C++ code (and consumers such as onnxruntime) that refers to the literal onnx namespace — rather than the ONNX_NAMESPACE macro — resolves to the onnx-light namespace. The standard onnx package lives in namespace onnx; onnx-light uses onnx_light (via ONNX_LIGHT_NAMESPACE), so this alias keeps onnx-light a true drop-in. It is only introduced when the onnx-light namespace differs from onnx.

namespace onnx_kernels

namespace kernel

class Attention : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Reference implementation of ai.onnx::Attention (v23 / v24).

Computes scaled dot-product attention. The baseline computation is

Y = Softmax((Q @ K^T) * scale + attn_bias, axis=-1) @ V

where attn_bias is derived from the optional attn_mask input (and/or the is_causal attribute) as described below.

Both rank-4 layouts (batch, num_heads, seq, head_size) and rank-3 “fused” layouts (batch, seq, num_heads * head_size) are accepted. When rank-3 inputs are used the q_num_heads and kv_num_heads attributes must be set and Y is returned with the same rank-3 layout as the inputs.

Supported features (mirroring the upstream operator):

• scale attribute — when unset the implementation uses 1 / sqrt(head_size).

• is_causal attribute — applies an upper-triangular -inf mask aligned to the upper-left corner of the (q_seq_len, total_kv_seq_len) score matrix (accounting for any cached KV).

• softcap attribute — when > 0, the pre-softmax scores are scaled by softcap * tanh(s / softcap).

• qk_matmul_output_mode attribute (0/1/2/3) — selects which stage of the QK pipeline is exposed in the auxiliary qk_matmul output (0: raw (Q @ K^T) * scale; 1: with bias; 2: after softcap; 3: after softmax).

• Grouped Query Attention (q_num_heads multiple of kv_num_heads).

• Optional attn_mask input: FLOAT (rank 2-4, broadcasted to (batch, q_num_heads, q_seq_len, total_kv_seq_len) and added as a bias) or BOOL (rank 2-4, true means “attend”; false becomes -inf). When the mask’s last dimension is shorter than total_kv_seq_len it is padded with -inf / false.

• Optional past_key/past_value rank-4 FLOAT inputs that are concatenated along the sequence axis with K/V to form the present_key/present_value outputs.

• Optional nonpad_kv_seqlen INT64 input (v24): a 1-D tensor of length batch_size giving the number of non-padding key/value positions per batch. Positions j >= nonpad_kv_seqlen[b] are masked out with -inf.

Not modeled: the softmax_precision attribute.

Public Functions

Tensor operator()(const Tensor &Q, const Tensor &K, const Tensor &V) const

Default-attribute overload. scale = 1 / sqrt(head_size), no mask, no causal, no softcap, no past KV.

Tensor operator()(const Tensor &Q, const Tensor &K, const Tensor &V, float scale) const

Explicit scale overload, no mask.

Tensor operator()(const Tensor &Q, const Tensor &K, const Tensor &V, float scale, const Tensor &attn_mask) const

Explicit scale plus an optional FLOAT attention bias whose shape must broadcast to (batch, q_num_heads, q_seq_len, kv_seq_len). A default-constructed Tensor omits the mask.

void operator()(const Tensor &Q, const Tensor &K, const Tensor &V, float scale, const Tensor *attn_mask, Tensor &output) const

In-place overload writing into a caller-allocated output tensor. output must already be FLOAT and shape (batch, q_num_heads, q_seq_len, v_head_size).

Result operator()(const Tensor &Q, const Tensor &K, const Tensor &V, const Attributes &attrs, const Tensor *attn_mask = nullptr, const Tensor *past_key = nullptr, const Tensor *past_value = nullptr, const Tensor *nonpad_kv_seqlen = nullptr) const

Comprehensive overload covering every supported feature. attn_mask, past_key, past_value and nonpad_kv_seqlen may be nullptr to mean “absent”. nonpad_kv_seqlen is a 1-D INT64 tensor of length batch_size masking out trailing (padding) key/value positions.

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Attention computes a fresh output buffer from independent reads of Q, K, V and never aliases an input buffer.

struct Attributes

#include <include_nn_kernels.h>

Bundles every optional scalar/string attribute of the upstream ai.onnx::Attention operator. has_scale controls whether scale overrides the default 1 / sqrt(head_size). When the inputs are rank-3 (fused layout), q_num_heads and kv_num_heads must both be positive; otherwise both default to 0 (auto-detected from the rank-4 input shapes).

Public Members

bool has_scale = false

When true, scale overrides the default.

float scale = 0.0f

Multiplier applied to the Q @ K^T dot product before the optional bias and softmax. When has_scale is false the kernel uses 1 / sqrt(head_size) instead.

bool is_causal = false

When true, applies a causal upper-triangular -inf mask.

float softcap = 0.0f

When > 0, applies softcap * tanh(s / softcap).

int qk_matmul_output_mode = 0

0: raw; 1: + bias; 2: after softcap; 3: after softmax.

int64_t q_num_heads = 0

Required when inputs are rank-3.

int64_t kv_num_heads = 0

Required when inputs are rank-3.

struct Result

#include <include_nn_kernels.h>

Bundles every output produced by the kernel. Y is always populated; present_key / present_value mirror the upstream optional outputs and equal K / V (or their concatenation with the past tensors when those were supplied). qk_matmul_output reflects the chosen qk_matmul_output_mode and is always rank-4 FLOAT (batch, q_num_heads, q_seq_len, total_kv_seq_len).

Public Members

Tensor Y

Primary output (same rank as Q).

Tensor present_key

Rank-4 (batch, kv_num_heads, total_kv_seq_len, head_size).

Tensor present_value

Rank-4 (batch, kv_num_heads, total_kv_seq_len, v_head_size).

Tensor qk_matmul_output

Auxiliary tensor (mode-dependent).

class AveragePool : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

N-D average pooling on a FLOAT tensor laid out as (N, C, D1, ..., Dk). kernel_shape must have k entries; strides, pads and dilations (lengths k, 2 * k and k respectively) default to all-ones / all-zeros / all-ones when omitted. auto_pad defaults to NOTSET (use explicit pads); when set to SAME_UPPER, SAME_LOWER or VALID the pads argument must be empty and the begin/end padding is computed from the input shape.

Public Functions

Tensor operator()(const Tensor &x, const std::vector<int64_t> &kernel_shape, const std::vector<int64_t> &strides = {}, const std::vector<int64_t> &pads = {}, bool ceil_mode = false, bool count_include_pad = false, const std::vector<int64_t> &dilations = {}, const std::string &auto_pad = "NOTSET") const

All attributes explicit. strides may be empty (treated as all 1), pads may be empty (treated as all 0) and dilations may be empty (treated as all 1).

void operator()(const Tensor &x, const std::vector<int64_t> &kernel_shape, const std::vector<int64_t> &strides, const std::vector<int64_t> &pads, bool ceil_mode, bool count_include_pad, Tensor &output, const std::vector<int64_t> &dilations = {}, const std::string &auto_pad = "NOTSET") const

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output shape generally differs from the input shape, so the output buffer cannot in general alias the input buffer.

class BatchNormalization : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Inference-mode BatchNormalization on a FLOAT input laid out as (N, C, D1, ..., Dk) (any rank >= 2; rank 1 is also accepted with C treated as 1). All four extra inputs (scale, B, input_mean, input_var) must be 1-D FLOAT tensors of length C. The kernel implements the inference formula Y = (X - input_mean) / sqrt(input_var + epsilon) * scale + B using NumPy-style broadcasting along the channel axis. Training mode (training_mode = 1, opset 14+) is also supported via TrainingForward: it normalizes X with the per-channel batch statistics and returns the updated running mean / variance.

Public Functions

Tensor operator()(const Tensor &x, const Tensor &scale, const Tensor &bias, const Tensor &input_mean, const Tensor &input_var, float epsilon = 1e-5f) const

Returns the inference-mode primary output Y. epsilon defaults to 1e-5f, the upstream default.

void operator()(const Tensor &x, const Tensor &scale, const Tensor &bias, const Tensor &input_mean, const Tensor &input_var, Tensor &output, float epsilon = 1e-5f) const

std::tuple<Tensor, Tensor, Tensor> TrainingForward(const Tensor &x, const Tensor &scale, const Tensor &bias, const Tensor &input_mean, const Tensor &input_var, float epsilon = 1e-5f, float momentum = 0.9f) const

Returns the training-mode outputs (Y, running_mean, running_var) for training_mode = 1 (opset 14+). The per-channel batch mean and (population) var are computed over every axis except the channel axis, Y is normalized with those batch statistics, and the running estimates are updated as running = input * momentum + saved * (1 - momentum). momentum defaults to 0.9f and epsilon to 1e-5f, the upstream defaults.

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output Y has the same shape as X so the output buffer may alias the input X buffer.

class CausalConvWithState : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Reference implementation of ai.onnx::CausalConvWithState (v27).

Stateful causal 1D depthwise convolution used by Gated DeltaNet (Qwen3.5) and Mamba (Jamba, FalconMamba). Replaces the Concat + Conv + Slice pattern with a single fused operation that also supports an optional fused SiLU/Swish activation.

Tensors are rank-3:

• input: (batch_size, channels, length)

• weight: (channels, 1, k) (depthwise: group = channels)

• bias (optional): (channels)

• past_state (optional): (batch_size, channels, k - 1); when absent the carry is implicitly zero-padding.

• output: same shape as input

• present_state: (batch_size, channels, k - 1) containing the last k - 1 values of the padded sequence ([past_state, input]) along the causal axis.

Public Functions

std::pair<Tensor, Tensor> operator()(const Tensor &input, const Tensor &weight, const Tensor &bias, const Tensor &past_state, const Attributes &attrs) const

Returning overload. bias and past_state may be default-constructed (empty-shape) Tensor values to indicate the optional inputs are missing.

void operator()(const Tensor &input, const Tensor &weight, const Tensor *bias, const Tensor *past_state, const Attributes &attrs, Tensor &output, Tensor &present_state) const

In-place overload writing into caller-allocated output tensors. bias and past_state may be nullptr for the optional inputs.

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

struct Attributes

#include <include_nn_kernels.h>

Attributes carried by the CausalConvWithState operator.

Public Members

std::string activation = "none"

Fused activation. Allowed values: "none" (default), "silu", "swish" (alias for "silu").

class Col2Im : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Reference N-D Col2Im kernel restricted to FLOAT tensors. Implements the upstream ai.onnx::Col2Im operator (since opset 18), rearranging column blocks back into a multi-dimensional batched image.

Inputs:

• input — FLOAT rank-3 tensor of shape (N, C * product(block_shape), L) where L is the number of blocks enumerated in lexicographic order.

• image_shape — INT64 1-D tensor of length n_spatial giving the spatial dimensions (d_1, ..., d_n) of the output image.

• block_shape — INT64 1-D tensor of length n_spatial giving the (un-dilated) block shape.

Attributes mirror the ONNX schema:

• dilations (optional, defaults to all ones).

• pads (optional, defaults to all zeros, length 2 * n_spatial in [d1_begin, ..., dn_begin, d1_end, ..., dn_end] order).

• strides (optional, defaults to all ones).

Output: FLOAT rank-2 + n_spatial tensor of shape (N, C, image_shape[0], ..., image_shape[n_spatial - 1]). Overlapping block contributions are summed (matching PyTorch’s fold).

Public Functions

Tensor operator()(const Tensor &input, const Tensor &image_shape, const Tensor &block_shape, const Attributes &attrs) const

void operator()(const Tensor &input, const Tensor &image_shape, const Tensor &block_shape, const Attributes &attrs, Tensor &output) const

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output shape differs from input shape, so storage cannot be shared.

struct Attributes

#include <include_nn_kernels.h>

Attributes carried by the ONNX Col2Im operator.

Public Members

std::vector<int64_t> dilations

Defaults to all ones.

std::vector<int64_t> pads

Defaults to all zeros (length 2 * n_spatial).

std::vector<int64_t> strides

Defaults to all ones.

class Conv : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Reference N-D Conv kernel restricted to FLOAT tensors.

Implements the standard convolution defined by the upstream ai.onnx::Conv operator. Inputs are laid out as X.shape = (N, C, D1, ..., Dk), W.shape = (M, C/group, k1, ..., kk); output shape is (N, M, oD1, ..., oDk). Supports the kernel_shape, strides, pads, dilations, group and auto_pad attributes (NOTSET, SAME_UPPER, SAME_LOWER, VALID).

Public Functions

Tensor operator()(const Tensor &x, const Tensor &w, const Tensor &b, const Attributes &attrs) const

Returning overload. B may be a default-constructed (empty-shape) Tensor to indicate the optional bias is missing.

void operator()(const Tensor &x, const Tensor &w, const Tensor &b, const Attributes &attrs, Tensor &output) const

In-place overload writing into a caller-allocated FLOAT output.

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output shape generally differs from the input, so storage cannot be shared.

struct Attributes

#include <include_nn_kernels.h>

Attributes carried by the ONNX Conv operator.

Public Members

std::vector<int64_t> kernel_shape

Defaults to W.shape[2:].

std::vector<int64_t> strides

Defaults to all ones.

std::vector<int64_t> pads

Defaults to all zeros (length 2 * rank).

std::vector<int64_t> dilations

Defaults to all ones.

int64_t group = 1

Number of conv groups.

std::string auto_pad = "NOTSET"

NOTSET / SAME_UPPER / SAME_LOWER / VALID.

class ConvInteger : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Reference N-D ConvInteger kernel.

Inputs x and w are 8-bit integer tensors (int8 or uint8) and the output y is int32. The optional x_zero_point must be a scalar with the same dtype as x; the optional w_zero_point may be a scalar or a 1-D length-M tensor (per-output-channel) with the same dtype as w. Empty-shape Tensor indicates absence.

Public Functions

Tensor operator()(const Tensor &x, const Tensor &w, const Tensor &x_zero_point, const Tensor &w_zero_point, const Attributes &attrs) const

Returning overload.

void operator()(const Tensor &x, const Tensor &w, const Tensor &x_zero_point, const Tensor &w_zero_point, const Attributes &attrs, Tensor &output) const

In-place overload writing into a caller-allocated INT32 output.

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output dtype differs from input dtype, so storage cannot be shared.

struct Attributes

#include <include_nn_kernels.h>

Attributes carried by the ONNX ConvInteger operator.

Public Members

std::vector<int64_t> kernel_shape

std::vector<int64_t> strides

std::vector<int64_t> pads

std::vector<int64_t> dilations

int64_t group = 1

std::string auto_pad = "NOTSET"

class ConvTranspose : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Reference N-D ConvTranspose kernel restricted to FLOAT tensors.

Implements the upstream ai.onnx::ConvTranspose. Input layout is X.shape = (N, C, D1, ..., Dk) and W.shape = (C, M/group, k1, ..., kk). When output_shape is provided the per-axis padding is derived per the upstream spec; otherwise the output spatial dim is stride[i] * (iD[i] - 1) + output_padding[i] + ((k[i]-1)*dil[i]+1) - pads[start] - pads[end].

Public Functions

Tensor operator()(const Tensor &x, const Tensor &w, const Tensor &b, const Attributes &attrs) const

Returning overload.

void operator()(const Tensor &x, const Tensor &w, const Tensor &b, const Attributes &attrs, Tensor &output) const

In-place overload writing into a caller-allocated FLOAT output.

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output shape generally differs from the input, so storage cannot be shared.

struct Attributes

#include <include_nn_kernels.h>

Attributes carried by the ONNX ConvTranspose operator.

Public Members

std::vector<int64_t> kernel_shape

std::vector<int64_t> strides

std::vector<int64_t> pads

std::vector<int64_t> dilations

std::vector<int64_t> output_padding

std::vector<int64_t> output_shape

int64_t group = 1

std::string auto_pad = "NOTSET"

class DeformConv : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Reference 2-D DeformConv kernel restricted to FLOAT tensors.

Implements the deformable convolution defined in https://arxiv.org/abs/1703.06211 and https://arxiv.org/abs/1811.11168, matching the upstream ai.onnx::DeformConv operator (since opset 19; the opset 22 update only widens T to also accept bfloat16 — not modeled here). Only the rank-4 (2-D image) case X.shape = (N, C, H, W) is supported.

Attributes mirror the ONNX schema:

• kernel_shape (optional): inferred from W.shape[2:] when omitted.

• strides / dilations (optional, default all ones).

• pads (optional, default all zeros, length 2 * spatial_rank in [h_begin, w_begin, h_end, w_end] order).

• group (default 1) — channels grouped along C / oC.

• offset_group (default 1) — number of offset groups; C must be divisible by offset_group.

Inputs:

• x — FLOAT rank-4 (N, C, H, W).

• w — FLOAT rank-4 (oC, C/group, kH, kW).

• offset — FLOAT rank-4 (N, offset_group * kH * kW * 2, oH, oW); the inner 2 channels are stored (y, x).

• b — Optional 1-D FLOAT bias of length oC; an empty-shape Tensor indicates absence (treated as zeros).

• mask — Optional FLOAT rank-4 mask (N, offset_group * kH * kW, oH, oW); an empty-shape Tensor indicates absence (treated as ones).

Public Functions

Tensor operator()(const Tensor &x, const Tensor &w, const Tensor &offset, const Tensor &b, const Tensor &mask, const Attributes &attrs) const

void operator()(const Tensor &x, const Tensor &w, const Tensor &offset, const Tensor &b, const Tensor &mask, const Attributes &attrs, Tensor &output) const

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output shape differs from any input, so storage cannot be shared.

struct Attributes

#include <include_nn_kernels.h>

Attributes carried by the ONNX DeformConv operator. Defaults match the upstream schema (since opset 19).

Public Members

std::vector<int64_t> kernel_shape

Defaults to W.shape[2:].

std::vector<int64_t> strides

Defaults to all ones.

std::vector<int64_t> pads

Defaults to all zeros (length 2 * rank).

std::vector<int64_t> dilations

Defaults to all ones.

int64_t group = 1

Number of conv groups.

int64_t offset_group = 1

Number of offset groups.

class Dropout : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Reference implementation of Dropout (opset 12+ behavior).

Dropout takes an input tensor data and optional scalar ratio / training_mode controls. In inference mode (training_mode=false), output is a copy of data and mask (when requested) is all ones. In training mode, each element is kept with probability 1-ratio and scaled by 1/(1-ratio).

This backend test kernel currently supports FLOAT and DOUBLE tensors for data. ratio must be in [0, 1).

Public Functions

std::pair<Tensor, Tensor> operator()(const Tensor &data, float ratio = 0.5f, bool training_mode = false, int64_t seed = kNoSeed) const

Tensor operator()(const Tensor &data, float ratio, bool training_mode, Tensor &mask, int64_t seed = kNoSeed) const

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Public Static Attributes

static constexpr int64_t kNoSeed = -1

class Flatten : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Flattens an input tensor of any element type into a 2-D matrix. Given an input of shape (d_0, d_1, ..., d_n) and integer attribute axis, the output has shape (d_0 * d_1 * ... * d_(axis-1), d_axis * ... * d_n). axis defaults to 1 and may be negative (counted from the back). When axis == 0 the outer dimension is 1 and the inner dimension is the product of every input dimension. When axis == rank the outer dimension is the product of every input dimension and the inner dimension is 1.

The output reuses the input data buffer byte-for-byte (Flatten is a pure view operation), so every element type supported by :cpp:func:PackedByteSize is accepted. The output dtype always matches the input dtype.

Public Functions

Tensor operator()(const Tensor &input, int64_t axis = 1) const

Returns the 2-D output tensor. axis follows ONNX semantics.

void operator()(const Tensor &input, int64_t axis, Tensor &output) const

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output element count equals the input element count, so storage can be shared in-place.

class GlobalAveragePool : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Global average pooling on a FLOAT tensor laid out as (N, C, D1, ..., Dk). The output shape is (N, C, 1, 1, ..., 1) — each spatial dimension is reduced to 1 by computing the mean over all elements in that dimension.

Public Functions

Tensor operator()(const Tensor &x) const

Returns a FLOAT output tensor of shape (N, C, 1, 1, ..., 1).

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output shape differs from the input shape, so storage cannot be shared.

class GlobalLpPool : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Global Lp pooling on a FLOAT tensor laid out as (N, C, D1, ..., Dk). The output shape is (N, C, 1, 1, ..., 1) — each spatial dimension is reduced to 1 by computing the Lp norm over all elements in that dimension. The default value of p is 2 (L2 norm). When p == 1 this is L1 pooling; p == 2 (default) gives the RMS/L2 pooling defined by ONNX.

Public Functions

Tensor operator()(const Tensor &x, int64_t p = 2) const

Returns a FLOAT output tensor of shape (N, C, 1, 1, ..., 1). p is the Lp norm exponent (default 2).

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output shape differs from the input shape, so storage cannot be shared.

class GlobalMaxPool : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Global max pooling on a FLOAT tensor laid out as (N, C, D1, ..., Dk). The output shape is (N, C, 1, 1, ..., 1) — each spatial dimension is reduced to 1 by computing the maximum over all elements in that dimension.

Public Functions

Tensor operator()(const Tensor &x) const

Returns a FLOAT output tensor of shape (N, C, 1, 1, ..., 1).

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output shape differs from the input shape, so storage cannot be shared.

class GroupNormalization : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Reference implementation of GroupNormalization (opset 18, 21).

Groups the C channel dimension of a FLOAT input X of shape (N, C, D1, ..., Dk) (rank >= 2) into num_groups equal-sized groups, normalizes each (n, group) block over its channels and spatial dims, then applies a per-channel affine scale * x + bias:

Y = scale * (X - mean) / sqrt(var + epsilon) + bias

The kernel implements the opset-21 contract where scale and bias are 1-D FLOAT tensors of length C. num_groups must divide C.

Public Functions

Tensor operator()(const Tensor &x, const Tensor &scale, const Tensor &bias, int64_t num_groups, float epsilon = 1e-5f) const

Returns the output Y. epsilon defaults to 1e-5f, the upstream default.

void operator()(const Tensor &x, const Tensor &scale, const Tensor &bias, int64_t num_groups, Tensor &output, float epsilon = 1e-5f) const

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output Y has the same shape as X so the output buffer may alias the input X buffer.

class GRU : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Single-direction ("forward") one-layer GRU on FLOAT tensors using the default Sigmoid/Tanh activations. Implements the upstream ONNX GRU formula

z_t = sigmoid(X_t @ Wz^T + H_{t-1} @ Rz^T + Wbz + Rbz) r_t = sigmoid(X_t @ Wr^T + H_{t-1} @ Rr^T + Wbr + Rbr) h_t = tanh (X_t @ Wh^T + (r_t (.) H_{t-1}) @ Rh^T + Wbh + Rbh) (linear_before_reset == 0, the default) h_t = tanh (X_t @ Wh^T + r_t (.) (H_{t-1} @ Rh^T + Rbh) + Wbh) (linear_before_reset != 0) H_t = (1 - z_t) (.) h_t + z_t (.) H_{t-1}

for both layout=0 (X.shape = [seq_length, batch_size, input_size]; optional initial_h.shape = [1, batch_size, hidden_size]) and layout=1 (X.shape = [batch_size, seq_length, input_size]; optional initial_h.shape = [batch_size, 1, hidden_size]). W.shape = [1, 3 * hidden_size, input_size]; R.shape = [1, 3 * hidden_size, hidden_size]; optional B.shape = [1, 6 * hidden_size] ([Wb, Rb] each with 3 gate blocks in the ONNX gate order z, r, h); initial_h defaults to zeros. sequence_lens is not supported (every batch must share the same sequence length); activations, clip and non-forward direction are not supported.

The two outputs are produced together: for layout=0, Y has shape [seq_length, 1, batch_size, hidden_size] and is the concatenation of every per-time-step hidden state; Y_h has shape [1, batch_size, hidden_size] and equals the last time step of Y. For layout=1 the corresponding shapes are [batch_size, seq_length, 1, hidden_size] and [batch_size, 1, hidden_size].

Public Functions

std::pair<Tensor, Tensor> operator()(const Tensor &x, const Tensor &w, const Tensor &r, const Tensor &b = Tensor{}, const Tensor &initial_h = Tensor{}, int64_t linear_before_reset = 0, int64_t layout = 0) const

Returns the pair (Y, Y_h). b may be a default-constructed (empty-shape) Tensor to indicate the optional B input is missing; same convention for initial_h. linear_before_reset matches the ONNX attribute of the same name (default 0). layout matches the ONNX attribute of the same name (default 0).

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output shape generally differs from the input shape, so storage cannot in general be shared.

class InstanceNormalization : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Reference implementation of InstanceNormalization (opset 1, 6, 22).

Normalizes a FLOAT input X of shape (N, C, D1, ..., Dk) (rank >= 2) per-instance, per-channel:

Y = scale * (X - mean) / sqrt(var + epsilon) + B

where mean and var are computed over the spatial dimensions D1, ..., Dk for every (n, c) pair. scale and B are 1-D FLOAT tensors of length C.

Public Functions

Tensor operator()(const Tensor &x, const Tensor &scale, const Tensor &bias, float epsilon = 1e-5f) const

Returns the output Y. epsilon defaults to 1e-5f, the upstream default.

void operator()(const Tensor &x, const Tensor &scale, const Tensor &bias, Tensor &output, float epsilon = 1e-5f) const

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output Y has the same shape as X so the output buffer may alias the input X buffer.

class LayerNormalization : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Reference implementation of LayerNormalization (opset 17).

Normalizes a FLOAT input X of arbitrary rank r along the last r - axis dimensions, then applies a per-element affine transform using Scale and the optional B (both broadcastable to the normalized shape X.shape[axis:]):

Mean = ReduceMean(X, axes=[axis,...,r-1]) Var = ReduceMean((X - Mean)^2, axes=[axis,...,r-1]) InvStdDev = 1 / sqrt(Var + epsilon) Y = (X - Mean) * InvStdDev * Scale + B

In addition to the normalized output Y, the kernel also returns the per-row Mean and InvStdDev tensors. Their shape mirrors X with the trailing r - axis dimensions collapsed to 1.

Public Functions

std::tuple<Tensor, Tensor, Tensor> operator()(const Tensor &x, const Tensor &scale, const Tensor &b, int64_t axis = -1, float epsilon = 1e-5f) const

Returns the tuple (Y, Mean, InvStdDev). axis defaults to -1 and epsilon to 1e-5f to match the upstream defaults. b may be a default-constructed Tensor{} to indicate that the bias is absent.

void operator()(const Tensor &x, const Tensor &scale, const Tensor &b, Tensor &y, Tensor &mean, Tensor &inv_std_dev, int64_t axis = -1, float epsilon = 1e-5f) const

In-place overload. y must have the same shape and dtype as x. mean and inv_std_dev must have x’s outer shape (i.e. x.shape[:axis]) padded with 1s up to x’s rank.

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output Y has the same shape as X so the output buffer may alias the input X buffer.

class LinearAttention : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Reference LinearAttention kernel for FLOAT, FLOAT16 and BFLOAT16 tensors.

Implements the unified linear attention recurrence defined by the upstream ai.onnx::LinearAttention operator (since opset 27). Supports four update_rule variants: "linear", "gated", "delta", and "gated_delta". Inputs use 3D packed format (B, T, H*D); the recurrent state is 4D (B, H_kv, d_k, d_v). Group-query attention (GQA) is supported when q_num_heads is a positive multiple of kv_num_heads. Half-precision (FLOAT16/BFLOAT16) activations are promoted to FLOAT32 for the recurrence and the outputs are demoted back.

Public Functions

Tensor operator()(const Tensor &query, const Tensor &key, const Tensor &value) const

Default-attribute overload (update_rule=”linear”, scale=auto).

Result operator()(const Tensor &query, const Tensor &key, const Tensor &value, const Attributes &attrs, const Tensor *past_state = nullptr, const Tensor *decay = nullptr, const Tensor *beta = nullptr) const

Comprehensive overload covering every supported feature.

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

struct Attributes

#include <include_nn_kernels.h>

Attributes carried by the ONNX LinearAttention operator.

Public Members

std::string update_rule = "gated_delta"

Recurrence type.

bool has_scale = false

When true, scale overrides the default.

float scale = 0.0f

Output scaling factor; 0 means 1/sqrt(d_k).

int64_t q_num_heads = 0

Number of query heads (required).

int64_t kv_num_heads = 0

Number of key/value heads (required).

int64_t chunk_size = 64

Chunk size hint (not used in reference impl).

struct Result

#include <include_nn_kernels.h>

Bundles every output produced by the kernel.

Public Members

Tensor output

3D packed (B, T, H_q * d_v).

Tensor present_state

4D (B, H_kv, d_k, d_v).

class LpNormalization : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Reference implementation of LpNormalization.

Computes output = input / Lp_norm(input, axis) element-wise along the given axis (negative values count from the end). When the Lp norm along a slice is zero, the corresponding output slice is set to zero to avoid division by zero.

The kernel accepts FLOAT inputs of any rank >= 1. axis defaults to -1 (the last axis). p must be either 1 (L1 norm) or 2 (L2 norm) and defaults to 2. The output shape matches the input shape.

Public Functions

Tensor operator()(const Tensor &x, int64_t axis = -1, int64_t p = 2) const

Returns a FLOAT output tensor whose shape matches x.

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output shape matches input shape, so storage may alias the input buffer.

class LpPool : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

N-D Lp pooling on a FLOAT tensor laid out as (N, C, D1, ..., Dk). Implements the ONNX LpPool operator: for every output position, the kernel computes (sum |x_i|^p)^(1/p) over the input values inside the pooling window. Positions outside [0, in_dim) (i.e. in the padded region) contribute 0 to the sum (matching the upstream ONNX reference, where the input is zero-padded before pooling).

kernel_shape must have k entries; strides, pads and dilations (lengths k, 2 * k and k respectively) default to all-ones / all-zeros / all-ones when omitted. auto_pad defaults to NOTSET (use explicit pads); when set to SAME_UPPER, SAME_LOWER or VALID the pads argument must be empty and the begin/end padding is computed from the input shape. p defaults to 2 (L2 norm).

Public Functions

Tensor operator()(const Tensor &x, const std::vector<int64_t> &kernel_shape, const std::vector<int64_t> &strides = {}, const std::vector<int64_t> &pads = {}, int64_t p = 2, bool ceil_mode = false, const std::vector<int64_t> &dilations = {}, const std::string &auto_pad = "NOTSET") const

All attributes explicit. strides may be empty (treated as all 1), pads may be empty (treated as all 0) and dilations may be empty (treated as all 1).

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output shape generally differs from the input shape, so the output buffer cannot in general alias the input buffer.

class LRN : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Reference implementation of LRN (Local Response Normalization).

Implements

Y[n, c, d1, ..., dk] = X[n, c, d1, ..., dk] / (bias + alpha / size * square_sum[n, c, d1, ..., dk]) ** beta

where square_sum[n, c, d1, ..., dk] is the sum of squares over the channel window i in [max(0, c - floor((size-1)/2)), min(C-1, c + ceil((size-1)/2))].

The kernel accepts FLOAT inputs of rank >= 2 laid out as (N, C, D1, ..., Dk). size must be a strictly positive integer. alpha defaults to 0.0001, beta to 0.75 and bias to 1.0 (the ONNX defaults). The output shape matches the input shape.

Public Functions

Tensor operator()(const Tensor &x, int64_t size, float alpha = 0.0001f, float beta = 0.75f, float bias = 1.0f) const

Returns a FLOAT output tensor whose shape matches x.

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output shape matches input shape, so storage may alias the input buffer.

class LSTM : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Single-direction ("forward") one-layer LSTM on FLOAT tensors using the default Sigmoid/Tanh/Tanh activations. Implements the upstream ONNX LSTM formula

it = sigmoid(Xt @ Wi^T + Ht-1 @ Ri^T + Pi (.) Ct-1 + Wbi + Rbi) ft = sigmoid(Xt @ Wf^T + Ht-1 @ Rf^T + Pf (.) Ct-1 + Wbf + Rbf) ct = tanh (Xt @ Wc^T + Ht-1 @ Rc^T + Wbc + Rbc) Ct = ft (.) Ct-1 + it (.) ct ot = sigmoid(Xt @ Wo^T + Ht-1 @ Ro^T + Po (.) Ct + Wbo + Rbo) Ht = ot (.) tanh(Ct)

for W.shape = [1, 4 * hidden_size, input_size]; R.shape = [1, 4 * hidden_size, hidden_size]; optional B.shape = [1, 8 * hidden_size] ([Wb, Rb] each with 4 gate blocks in the ONNX gate order i, o, f, c); optional P.shape = [1, 3 * hidden_size] (peephole weights in gate order i, o, f). sequence_lens is not supported (every batch must share the same sequence length); activations, clip, input_forget and non-forward direction are not supported.

Both layout values defined by ONNX are supported:

• layout == 0 (default): X.shape = [seq_length, batch_size, input_size]; optional initial_h.shape = initial_c.shape = [1, batch_size, hidden_size]; Y.shape = [seq_length, 1, batch_size, hidden_size]; Y_h.shape = [1, batch_size, hidden_size].

• layout == 1: X.shape = [batch_size, seq_length, input_size]; optional initial_h.shape = initial_c.shape = [batch_size, 1, hidden_size]; Y.shape = [batch_size, seq_length, 1, hidden_size]; Y_h.shape = [batch_size, 1, hidden_size].

The two outputs are produced together: Y is the concatenation of every per-time-step hidden state and Y_h equals the last time step of Y. The optional third output Y_c is not produced by this overload.

Public Functions

std::pair<Tensor, Tensor> operator()(const Tensor &x, const Tensor &w, const Tensor &r, const Tensor &b = Tensor{}, const Tensor &initial_h = Tensor{}, const Tensor &initial_c = Tensor{}, const Tensor &p = Tensor{}, int64_t layout = 0) const

Returns the pair (Y, Y_h). b, initial_h, initial_c and p may each be a default-constructed (empty-shape) Tensor to indicate that the corresponding optional input is missing. layout selects between the two ONNX layouts (0 or 1); any other value is rejected.

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output shape generally differs from the input shape, so storage cannot in general be shared.

class MaxPool : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

N-D max pooling on a numeric tensor laid out as (N, C, D1, ..., Dk). Supported element types are FLOAT, DOUBLE, INT8 and UINT8; the output type matches the input type. Implements the ONNX MaxPool operator restricted to storage_order=0 (row-major). kernel_shape must have k entries; strides, pads and dilations (lengths k, 2 * k and k respectively) default to all-ones / all-zeros / all-ones when omitted. auto_pad defaults to NOTSET (use explicit pads); when set to SAME_UPPER, SAME_LOWER or VALID the pads argument must be empty and the begin/end padding is computed from the input shape.

The :cpp:func:operator() overload returns only the primary output Y. :cpp:func:WithIndices additionally produces the optional Indices tensor (int64 flat indices into the un-padded input buffer, mirroring the upstream ONNX semantics).

Public Functions

Tensor operator()(const Tensor &x, const std::vector<int64_t> &kernel_shape, const std::vector<int64_t> &strides = {}, const std::vector<int64_t> &pads = {}, bool ceil_mode = false, const std::vector<int64_t> &dilations = {}, int64_t storage_order = 0, const std::string &auto_pad = "NOTSET") const

Returns the primary output Y (tensor of the pooled values; same element type as x).

std::pair<Tensor, Tensor> WithIndices(const Tensor &x, const std::vector<int64_t> &kernel_shape, const std::vector<int64_t> &strides = {}, const std::vector<int64_t> &pads = {}, bool ceil_mode = false, const std::vector<int64_t> &dilations = {}, int64_t storage_order = 0, const std::string &auto_pad = "NOTSET") const

Returns (Y, Indices) where Indices is an int64 tensor with the same shape as Y containing the flat indices into the un-padded input (N, C, D1, ..., Dk) buffer. storage_order selects the flattening of the selected spatial coordinate: 0 (row major) or 1 (column major).

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output shape generally differs from the input shape.

class MaxRoiPool : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

2-D MaxRoiPool kernel restricted to FLOAT inputs/outputs. Implements the ONNX MaxRoiPool operator: for each RoI [batch_id, x1, y1, x2, y2] the kernel scales the corner coordinates by spatial_scale, rounds them to the nearest integer feature-map cell, divides the resulting region into a pooled_shape[0] x pooled_shape[1] grid of bins, and produces one max value per (channel, bin) pair. Empty bins (e.g. when the scaled RoI is smaller than the requested pooled output) emit 0 — mirroring the upstream reference implementation in onnx/reference/ops/op_max_roi_pool.py.

Public Functions

Tensor operator()(const Tensor &x, const Tensor &rois, const Attributes &attrs) const

Returns the pooled output Y of shape (num_rois, C, pooled_shape[0], pooled_shape[1]).

void operator()(const Tensor &x, const Tensor &rois, const Attributes &attrs, Tensor &output) const

In-place variant that writes into a pre-allocated output tensor of the expected shape and FLOAT dtype.

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output element layout (num_rois, C, pooled_shape[0], pooled_shape[1]) differs from any input, so storage cannot be shared.

struct Attributes

#include <include_nn_kernels.h>

Attributes carried by the ONNX MaxRoiPool operator.

Public Members

std::vector<int64_t> pooled_shape

ROI pool output shape (height, width). Required; must contain exactly two positive entries.

float spatial_scale = 1.0f

Multiplicative spatial scale factor (default 1.0).

class MaxUnpool : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

N-D MaxUnpool on a FLOAT tensor with int64 indices. Implements the ONNX MaxUnpool operator: each value of x is scattered into a zero-initialized output buffer at the linear index given by indices (interpreted as the flat (N, C, D1, ..., Dk) offset of the output).

Two overloads are provided: one that derives the output shape from the kernel_shape, strides and pads attributes, and one that takes an explicit rank-1 int64 output_shape tensor (if specified, the pads attribute is ignored, matching the ONNX spec).

Public Functions

Tensor operator()(const Tensor &x, const Tensor &indices, const std::vector<int64_t> &kernel_shape, const std::vector<int64_t> &strides = {}, const std::vector<int64_t> &pads = {}) const

Two-input form: derives the output shape from kernel_shape, strides and pads (defaults: 1 and 0).

Tensor operator()(const Tensor &x, const Tensor &indices, const Tensor &output_shape, const std::vector<int64_t> &kernel_shape, const std::vector<int64_t> &strides = {}, const std::vector<int64_t> &pads = {}) const

Three-input form: output_shape is a rank-1 int64 tensor giving the full output shape (N, C, D1, ..., Dk).

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output shape differs from input shape.

class MeanVarianceNormalization : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Reference implementation of MeanVarianceNormalization (opset 9, 13).

Normalizes each element as Y = (X - mean(X, axes)) / sqrt(mean((X - mean(X, axes))^2, axes)). The default axes are [0, 2, 3].

Public Functions

Tensor operator()(const Tensor &x, const std::vector<int64_t> &axes = {0, 2, 3}) const

void operator()(const Tensor &x, Tensor &output, const std::vector<int64_t> &axes = {0, 2, 3}) const

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output Y has the same shape as X so the output buffer may alias the input X buffer.

class RMSNormalization : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Reference implementation of RMSNormalization (opset 23).

Normalizes a FLOAT input X of arbitrary rank r by dividing it by the root-mean-square computed over the last r - axis dimensions, then multiplies the normalized result by scale (with unidirectional broadcasting):

Y = X / sqrt(mean(X * X, axes=normalized_axes) + epsilon) * scale

where normalized_axes = [axis, ..., r - 1]. scale must be broadcastable to the normalized shape (i.e. X.shape[axis:]).

Public Functions

Tensor operator()(const Tensor &x, const Tensor &scale, int64_t axis = -1, float epsilon = 1e-5f) const

Returns the output Y. axis defaults to -1 and epsilon to 1e-5f to match the upstream defaults.

void operator()(const Tensor &x, const Tensor &scale, Tensor &output, int64_t axis = -1, float epsilon = 1e-5f) const

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output Y has the same shape as X so the output buffer may alias the input X buffer.

class RNN : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Single-direction ("forward") one-layer RNN on FLOAT tensors using the Tanh activation. Implements the upstream ONNX RNN formula

H_t = tanh(X_t @ W^T + W_b + H_{t-1} @ R^T + R_b)

for layout=0, with X.shape = [seq_length, batch_size, input_size]; W.shape = [1, hidden_size, input_size]; R.shape = [1, hidden_size, hidden_size]; optional B.shape = [1, 2 * hidden_size] ([Wb, Rb]); optional initial_h.shape = [1, batch_size, hidden_size], defaulting to zeros. layout=1 (batch-major) is also supported: the kernel transposes X/initial_h on entry and the outputs on exit so the core time-major loop is unchanged. The activations attribute, if present, must be either empty or the single value "Tanh"; direction must be "forward" (the default); sequence_lens is not supported (every batch must share the same sequence length); clip is not supported.

The two outputs are produced together: Y has shape [seq_length, 1, batch_size, hidden_size] and is the concatenation of every per-time-step hidden state; Y_h has shape [1, batch_size, hidden_size] and equals the last time step of Y.

Public Functions

std::pair<Tensor, Tensor> operator()(const Tensor &x, const Tensor &w, const Tensor &r, const Tensor &b = Tensor{}, const Tensor &initial_h = Tensor{}, int64_t layout = 0) const

Returns the pair (Y, Y_h). b may be a default-constructed (empty-shape) Tensor to indicate the optional B input is missing; same convention for initial_h. layout may be 0 (time-major, the default) or 1 (batch-major); for layout=1 the kernel transposes X/initial_h on entry and the outputs on exit so the core time-major loop is unchanged.

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output shape generally differs from the input shape, so storage cannot in general be shared.

class RotaryEmbedding : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_nn_kernels.h>

Reference implementation of ai.onnx::RotaryEmbedding (v23).

Applies the rotary positional embedding (RoPE) defined in https://arxiv.org/pdf/2104.09864 to the input tensor X. The kernel accepts both the rank-3 layout (batch_size, sequence_length, hidden_size) (in which case num_heads must be provided and hidden_size = num_heads * head_size) and the rank-4 layout (batch_size, num_heads, sequence_length, head_size). The output has the same shape and dtype as X.

cos_cache and sin_cache are either rank-2 (max_position_id_plus_1, head_size / 2) tables indexed by the optional position_ids input or rank-3 (batch_size, sequence_length, head_size / 2) tensors used directly when position_ids is omitted. When rotary_embedding_dim is non-zero the last dimension of the caches must equal rotary_embedding_dim / 2 and rotation is applied only to the first rotary_embedding_dim channels of each head; the remaining channels are passed through unchanged.

When interleaved is true, even/odd channels of each head are paired instead of the (default) splitting into two contiguous halves.

Public Functions

Tensor operator()(const Tensor &X, const Tensor &cos_cache, const Tensor &sin_cache, const Tensor &position_ids, const Attributes &attrs) const

Returning overload. position_ids may be a default-constructed (empty-shape) Tensor to indicate the optional input is missing.

void operator()(const Tensor &X, const Tensor &cos_cache, const Tensor &sin_cache, const Tensor *position_ids, const Attributes &attrs, Tensor &output) const

In-place overload writing into a caller-allocated output of the same shape and dtype as X.

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

Output shape matches input shape; storage could theoretically alias X but the in-place overload above does not exploit it.

struct Attributes

#include <include_nn_kernels.h>

Attributes carried by the ONNX RotaryEmbedding operator. All defaults match the upstream schema.

Public Members

bool interleaved = false

Pair (0,1),(2,3),… instead of (0,h/2),(1,h/2+1),…

int64_t rotary_embedding_dim = 0

0 means “rotate all of head_size”.

int64_t num_heads = 0

Required when X is rank-3.