#[non_exhaustive]pub enum Op {
Show 187 variants
Input {
name: String,
},
Param {
name: String,
},
Constant {
data: Vec<u8>,
},
Activation(Activation),
Cast {
to: DType,
},
StopGradient,
Quantize {
axis: Option<usize>,
scales: Vec<f32>,
zero_points: Vec<i32>,
},
Dequantize {
axis: Option<usize>,
scales: Vec<f32>,
zero_points: Vec<i32>,
},
FakeQuantize {
bits: u8,
axis: Option<usize>,
ste: SteKind,
scale_mode: ScaleMode,
},
FakeQuantizeLSQ {
bits: u8,
axis: Option<usize>,
},
FakeQuantizeLSQBackwardX {
bits: u8,
axis: Option<usize>,
},
FakeQuantizeLSQBackwardScale {
bits: u8,
axis: Option<usize>,
},
Binary(BinaryOp),
Compare(CmpOp),
Where,
Fma,
ElementwiseRegion {
chain: Vec<ChainStep>,
num_inputs: u32,
scalar_input_mask: u32,
input_modulus: [u32; 16],
prologue: RegionPrologue,
prologue_input: u32,
},
TransformRegion {
steps: Vec<TransformStep>,
num_inputs: u32,
},
BatchElementwiseRegion {
chain: Vec<ChainStep>,
num_batch_inputs: u32,
scalar_input_mask: u32,
input_modulus: [u32; 16],
prologue: RegionPrologue,
prologue_input: u32,
},
MatMul,
DotGeneral {
lhs_contracting: Vec<usize>,
rhs_contracting: Vec<usize>,
lhs_batch: Vec<usize>,
rhs_batch: Vec<usize>,
},
BatchedDenseSolve,
DenseSolve,
Cholesky,
TriangularSolve {
lower: bool,
transpose: bool,
},
Det,
LogDet,
Sort {
axis: usize,
descending: bool,
},
ArgSort {
axis: usize,
descending: bool,
},
Svd {
part: SvdPart,
},
Qr {
part: QrPart,
},
LayerNorm {
axis: i32,
eps: f32,
},
GroupNorm {
num_groups: usize,
eps: f32,
},
LayerNorm2d {
eps: f32,
},
ResizeNearest2x,
Interpolate3d {
size: Vec<usize>,
},
RmsNorm {
axis: i32,
eps: f32,
},
BatchNormInference {
eps: f32,
},
Attention {
num_heads: usize,
head_dim: usize,
v_head_dim: Option<usize>,
mask_kind: MaskKind,
score_scale: Option<f32>,
attn_logit_softcap: Option<f32>,
},
Rope {
head_dim: usize,
n_rot: usize,
style: RopeStyle,
},
AxialRope2d {
end_x: usize,
end_y: usize,
head_dim: usize,
num_heads: usize,
theta: f32,
repeat_factor: usize,
},
Reshape {
new_shape: Vec<i64>,
},
Transpose {
perm: Vec<usize>,
},
Narrow {
axis: usize,
start: usize,
len: usize,
},
Slice {
axis: usize,
start: usize,
len: usize,
step: i64,
},
Roll {
shifts: Vec<i64>,
dims: Vec<usize>,
},
Concat {
axis: usize,
},
KvAppend {
axis: usize,
pos: usize,
},
Expand {
target_shape: Vec<i64>,
},
Gather {
axis: usize,
},
Reverse {
axes: Vec<usize>,
},
Pad {
pads: Vec<[usize; 2]>,
mode: PadMode,
},
Clamp {
min: f32,
max: f32,
},
Tile {
reps: Vec<usize>,
},
Trilu {
upper: bool,
diagonal: i64,
},
Reduce {
op: ReduceOp,
axes: Vec<usize>,
keep_dim: bool,
},
Histogram {
bins: usize,
min: f32,
max: f32,
},
SelectiveScan {
state_size: usize,
},
GatedDeltaNet {
state_size: usize,
carry_state: bool,
gate_per_channel: bool,
},
GatedDeltaNetBackward {
state_size: usize,
carry_state: bool,
gate_per_channel: bool,
},
Lstm {
hidden_size: usize,
num_layers: usize,
bidirectional: bool,
carry: bool,
},
Gru {
hidden_size: usize,
num_layers: usize,
bidirectional: bool,
carry: bool,
},
Rnn {
hidden_size: usize,
num_layers: usize,
bidirectional: bool,
carry: bool,
relu: bool,
},
Mamba2 {
head_dim: usize,
state_size: usize,
},
DequantMatMul {
scheme: QuantScheme,
},
SynthMatMul {
kind: SynthKind,
},
SynthMatMulBackward {
kind: SynthKind,
wrt: SynthBwdWrt,
},
SynthReconstruct {
kind: SynthKind,
},
SplineActivation {
num_basis: u32,
grid_min: f32,
grid_max: f32,
},
SplineActivationBackwardX {
num_basis: u32,
grid_min: f32,
grid_max: f32,
},
SplineActivationBackwardCoeff {
num_basis: u32,
grid_min: f32,
grid_max: f32,
},
QMatMul {
x_zp: i32,
w_zp: i32,
out_zp: i32,
mult: f32,
},
QConv2d {
kernel_size: Vec<usize>,
stride: Vec<usize>,
padding: Vec<usize>,
dilation: Vec<usize>,
groups: usize,
x_zp: i32,
w_zp: i32,
out_zp: i32,
mult: f32,
},
ScaledMatMul {
lhs_format: ScaledFormat,
rhs_format: ScaledFormat,
scale_layout: ScaleLayout,
has_bias: bool,
},
ScaledQuantize {
format: ScaledFormat,
scale_layout: ScaleLayout,
},
ScaledQuantScale {
format: ScaledFormat,
scale_layout: ScaleLayout,
},
ScaledDequantize {
format: ScaledFormat,
scale_layout: ScaleLayout,
},
LoraMatMul {
scale: f32,
},
PartitionedConv {
block: usize,
},
Sample {
top_k: usize,
top_p: f32,
temperature: f32,
seed: u64,
},
RngNormal {
mean: f32,
scale: f32,
key: u64,
op_seed: Option<f32>,
},
RngUniform {
low: f32,
high: f32,
key: u64,
op_seed: Option<f32>,
},
Cumsum {
axis: i32,
exclusive: bool,
},
CumProd {
axis: i32,
exclusive: bool,
},
CumMax {
axis: i32,
exclusive: bool,
},
ArgMax {
axis: usize,
keep_dim: bool,
},
ArgMin {
axis: usize,
keep_dim: bool,
},
Softmax {
axis: i32,
},
TopK {
k: usize,
},
GroupedMatMul,
DequantGroupedMatMul {
scheme: QuantScheme,
},
DequantGroupedMatMulMlx {
scheme: QuantScheme,
},
DequantMoEWeights {
scheme: QuantScheme,
},
ScaledGroupedMatMul {
lhs_format: ScaledFormat,
rhs_format: ScaledFormat,
scale_layout: ScaleLayout,
has_bias: bool,
},
ScatterAdd {
axis: usize,
},
ScatterNd {
reduction: ScatterNdReduction,
},
ScatterElements {
axis: i32,
reduction: ScatterNdReduction,
},
GatherNd {
batch_dims: i32,
},
GatherElements {
axis: i32,
},
Conv {
kernel_size: Vec<usize>,
stride: Vec<usize>,
padding: Vec<usize>,
dilation: Vec<usize>,
groups: usize,
},
Im2Col {
kernel_size: Vec<usize>,
stride: Vec<usize>,
padding: Vec<usize>,
dilation: Vec<usize>,
},
ConvTranspose2d {
kernel_size: Vec<usize>,
stride: Vec<usize>,
padding: Vec<usize>,
dilation: Vec<usize>,
output_padding: Vec<usize>,
groups: usize,
},
Conv3d {
stride: [usize; 3],
padding: [usize; 3],
dilation: [usize; 3],
groups: usize,
},
ConvTranspose3d {
stride: [usize; 3],
padding: [usize; 3],
dilation: [usize; 3],
output_padding: [usize; 3],
groups: usize,
},
Pool {
kind: ReduceOp,
kernel_size: Vec<usize>,
stride: Vec<usize>,
padding: Vec<usize>,
},
ReluBackward,
ComplexNormSq,
Conjugate,
ComplexNormSqBackward,
LayerNormBackwardInput {
axis: i32,
eps: f32,
},
LayerNormBackwardGamma {
axis: i32,
eps: f32,
},
RmsNormBackwardInput {
axis: i32,
eps: f32,
},
RmsNormBackwardGamma {
axis: i32,
eps: f32,
},
RmsNormBackwardBeta {
axis: i32,
eps: f32,
},
RopeBackward {
head_dim: usize,
n_rot: usize,
style: RopeStyle,
},
GroupNormBackwardInput {
num_groups: usize,
eps: f32,
},
GroupNormBackwardGamma {
num_groups: usize,
eps: f32,
},
GroupNormBackwardBeta {
num_groups: usize,
eps: f32,
},
BatchNormInferenceBackwardInput {
eps: f32,
},
BatchNormInferenceBackwardGamma {
eps: f32,
},
BatchNormInferenceBackwardBeta,
CumsumBackward {
axis: i32,
exclusive: bool,
},
GatherBackward {
axis: i32,
},
ActivationBackward {
kind: Activation,
},
FakeQuantizeBackward {
bits: u8,
axis: Option<usize>,
ste: SteKind,
},
MaxPool2dBackward {
kernel_size: Vec<usize>,
stride: Vec<usize>,
padding: Vec<usize>,
},
Conv2dBackwardInput {
kernel_size: Vec<usize>,
stride: Vec<usize>,
padding: Vec<usize>,
dilation: Vec<usize>,
groups: usize,
},
Conv2dBackwardWeight {
kernel_size: Vec<usize>,
stride: Vec<usize>,
padding: Vec<usize>,
dilation: Vec<usize>,
groups: usize,
},
MaxPool3dBackward {
kernel_size: Vec<usize>,
stride: Vec<usize>,
padding: Vec<usize>,
},
Conv3dBackwardInput {
kernel_size: Vec<usize>,
stride: Vec<usize>,
padding: Vec<usize>,
dilation: Vec<usize>,
groups: usize,
},
Conv3dBackwardWeight {
kernel_size: Vec<usize>,
stride: Vec<usize>,
padding: Vec<usize>,
dilation: Vec<usize>,
groups: usize,
},
SoftmaxCrossEntropy,
SoftmaxCrossEntropyWithLogits,
SoftmaxCrossEntropyBackward,
AttentionBackward {
num_heads: usize,
head_dim: usize,
mask_kind: MaskKind,
wrt: AttentionBwdWrt,
},
AttentionBackwardAll {
num_heads: usize,
head_dim: usize,
mask_kind: MaskKind,
},
FusedMatMulBiasAct {
activation: Option<Activation>,
},
FusedMatMulResidual,
FusedConvBiasAct {
kernel_size: Vec<usize>,
stride: Vec<usize>,
padding: Vec<usize>,
dilation: Vec<usize>,
groups: usize,
activation: Option<Activation>,
has_residual: bool,
},
FusedResidualLN {
has_bias: bool,
eps: f32,
},
FusedResidualRmsNorm {
has_bias: bool,
eps: f32,
},
FusedSwiGLU {
cast_to: Option<DType>,
gate_first: bool,
},
FusedTransformerLayer {
num_heads: usize,
head_dim: usize,
intermediate_size: usize,
eps1: f32,
eps2: f32,
activation: Activation,
has_bias: bool,
},
FusedAttentionBlock {
num_heads: usize,
head_dim: usize,
has_bias: bool,
has_rope: bool,
},
If {
then_branch: Box<Graph>,
else_branch: Box<Graph>,
},
While {
cond: Box<Graph>,
body: Box<Graph>,
max_iterations: Option<usize>,
},
Scan {
body: Box<Graph>,
length: u32,
save_trajectory: bool,
num_bcast: u32,
num_xs: u32,
num_checkpoints: u32,
},
ScanBackward {
body_vjp: Box<Graph>,
length: u32,
save_trajectory: bool,
num_xs: u32,
num_checkpoints: u32,
forward_body: Option<Box<Graph>>,
},
ScanBackwardXs {
body_vjp: Box<Graph>,
length: u32,
save_trajectory: bool,
num_xs: u32,
xs_idx: u32,
num_checkpoints: u32,
forward_body: Option<Box<Graph>>,
},
GaussianSplatRender {
width: u32,
height: u32,
tile_size: u32,
radius_scale: f32,
alpha_cutoff: f32,
max_splat_steps: u32,
transmittance_threshold: f32,
max_list_entries: u32,
},
GaussianSplatRenderBackward {
width: u32,
height: u32,
tile_size: u32,
radius_scale: f32,
alpha_cutoff: f32,
max_splat_steps: u32,
transmittance_threshold: f32,
max_list_entries: u32,
loss_grad_clip: f32,
sh_band: u32,
max_anisotropy: f32,
},
GaussianSplatPrepare {
width: u32,
height: u32,
tile_size: u32,
radius_scale: f32,
alpha_cutoff: f32,
max_splat_steps: u32,
transmittance_threshold: f32,
max_list_entries: u32,
},
GaussianSplatRasterize {
width: u32,
height: u32,
tile_size: u32,
alpha_cutoff: f32,
max_splat_steps: u32,
transmittance_threshold: f32,
max_list_entries: u32,
},
Custom {
name: String,
num_inputs: u32,
attrs: Vec<u8>,
},
Fft {
inverse: bool,
norm: FftNorm,
},
FftQ {
inverse: bool,
norm: FftNorm,
scale: FftQScale,
},
FftButterflyStage {
stage: u32,
n_fft: u32,
},
LogMel,
LogMelBackward,
WelchPeaks {
k: usize,
n_segments: usize,
},
CustomFn {
fwd_body: Box<Graph>,
vjp_body: Option<Box<Graph>>,
jvp_body: Option<Box<Graph>>,
num_inputs: u32,
},
BiMap,
ReEig {
eps: f32,
},
LogEig {
eps: f32,
},
SpdBatchNorm {
eps: f32,
},
SpdKarcherMean {
iters: u32,
tol: f32,
},
SpdKarcherMeanWeighted {
iters: u32,
tol: f32,
},
SpdLogMap,
SpdExpMap,
SpdParallelTransport,
SpdMatrixFnBatch {
kind: SpdMatFn,
},
ReEigBackward {
eps: f32,
},
LogEigBackward {
eps: f32,
},
SpdBatchNormBackwardX {
eps: f32,
},
SpdBatchNormBackwardG {
eps: f32,
},
SpdLogMapBackward,
SpdExpMapBackward,
SpdParallelTransportBackward,
SpdMatrixFnBatchBackward {
kind: SpdMatFn,
},
Eigh,
EighBackward,
EighBatch,
EighBatchBackward,
AdaLayerNorm {
norm: AdaNormKind,
eps: f32,
},
GatedResidual,
AdaLayerNormBackward {
norm: AdaNormKind,
eps: f32,
},
GatedResidualBackward,
}Expand description
- Matmul/Conv are BLAS-dispatched and form fusion boundaries
- Reductions are fusion roots (drive the loop iteration)
Variants (Non-exhaustive)§
This enum is marked as non-exhaustive
Input
Model input with a name (shape on the Node).
Param
Model parameter (weight/bias) with a name.
Constant
Constant tensor embedded in the graph.
Activation(Activation)
Unary activation: one input, same shape output.
Cast
Cast to a different dtype.
StopGradient
Stop-gradient (a.k.a. detach / lax.stop_gradient). Forward is
identity; the reverse-mode autodiff rule returns no gradient
contribution for the input. Single input, same shape & dtype
output. Used to build a Gradient-Reverse-Layer with identity
forward semantics in user code (see maet-rs dat_loss).
Quantize
INT8 quantization. Input f32; output i8 same shape.
q[i] = saturate_i8(round(x[i] / scale[c]) + zero_point[c])
where c selects the per-channel scale/zp when axis = Some(d)
(c = idx[d]), or always uses index 0 when axis = None
(per-tensor). The scales / zero_points payload length must
match 1 for per-tensor and input.dim(d) for per-channel.
Static — typically produced at calibration time and baked
into the loaded model. Use Op::Dequantize for the inverse.
Dequantize
INT8 dequantization (inverse of Op::Quantize). Input i8;
output f32 same shape.
x[i] = (q[i] - zero_point[c]) · scale[c]
where c is selected by axis exactly as in Op::Quantize.
FakeQuantize
“Fake-quantize” op for quantization-aware training (QAT).
Input f32; output f32 same shape. Forward computes a per-axis
(or per-tensor when axis = None) max-abs scale on the fly:
s[c] = max(|x[..., c, ...]|) / q_max(bits)
then quantizes-then-dequantizes:
out[i] = clamp(round(x[i] / s[c]), -q_max, q_max) * s[c]
where q_max is 127 for bits=8, 7 for bits=4, 1 for
bits=2 (ternary). Symmetric only — zero-point is always 0.
The point of this op is to make the SGD optimizer “see” the
deployment-time rounding during training. Backward is the
straight-through estimator (STE): the gradient passes
through (variant chosen by ste), ignoring the discontinuity
at the round. Without STE the rounding would have zero
gradient almost everywhere and learning would stop.
Inserted by the trainer on conv / FC weight tensors when
--qat is on; the existing Op::Quantize / packing path at
the end of training still handles the deployment-side
conversion to i8/i4/i2 codes.
FakeQuantizeLSQ
Learned Step Size Quantization (LSQ; Esser et al. 2020,
arXiv:1902.08153). Like FakeQuantize but the per-channel
scale is a learned parameter, passed as the second input.
Forward is identical to FakeQuantize with a fixed scale:
out[i] = clamp(round(x[i]/s[c]), -q_max, q_max) * s[c]
Backward computes both dx (STE) and dscale[c] via the
closed-form gradient:
dscale[c] = sum_i ψ(x[i]/s[c]) · upstream[i]
where ψ(z) = -z + round(z) if |z| ≤ q_max else
sign(z) · q_max. Routinely beats per-batch and EMA at
tight bit widths (i2 / i3).
Inputs: [x, scale]. scale is [chan_dim] f32 (matches
axis); for axis = None it’s [1].
FakeQuantizeLSQBackwardX
Backward pass for Op::FakeQuantizeLSQ. Computes BOTH the
gradient w.r.t. x (STE) and the gradient w.r.t. scale
(closed-form). Output shape matches x; the scale gradient
is reduced separately by LsqScaleGradient.
Inputs: [x, scale, dy]. Output: dx, same shape as x.
FakeQuantizeLSQBackwardScale
Companion to FakeQuantizeLSQBackwardX: computes the
[chan_dim] per-channel scale gradient. Inputs [x, scale, dy].
Output shape matches scale.
Binary(BinaryOp)
Binary op with broadcasting: two inputs, output shape is broadcast result.
Compare(CmpOp)
Element-wise comparison: two inputs, Bool output.
Where
Select elements: cond (Bool), on_true, on_false → output.
Fma
Fused multiply-add with a SINGLE rounding: a*b + c. Distinct from a
mul then add (two roundings) — the single rounding is what makes
error-free transforms (TwoProduct) and compensated/double-word
arithmetic possible. Inputs: a, b, c.
The native (single-rounding) kernel is elementwise on equal-shaped
operands — broadcast a scalar to the common shape first. Backends
without a native FMA fall back to mul+add via LowerFma (two
roundings, but broadcasts), so this only constrains the precise path.
ElementwiseRegion
Fused element-wise region (PLAN L2). Holds an N-step chain of
element-wise operations. Inputs are referenced by index 0..num_inputs;
each step’s result can be referenced by later steps via
ChainOperand::Step(idx). The output is the last step’s result.
Emitted by MarkElementwiseRegions in rlx-opt from chains of
Activation/Cast/Binary/Compare/Where ops with single-consumer
intermediates and broadcast-compatible shapes. Backends that
don’t have a region kernel can decompose back to the original
chain via unfuse::unfuse_elementwise_regions.
scalar_input_mask is a per-input bitfield (bit i set ⇒
input i is a scalar broadcast — has shape [1]). Kept as a
fast-path indicator that lets kernels skip the modulo entirely
when they detect a scalar.
input_modulus[i] is the per-input element count, used by
kernels to compute arena[input_offs[i] + (gid % input_modulus[i])]
— the trailing-shape broadcast pattern. 0 means “no broadcast”
(input matches the output element count; kernel reads gid
directly). 1 means scalar; any other value means the input
has fewer elements than the output and they tile by modulo.
The encoder only allows broadcasts where out_elems % in_elems == 0 so the modulo divides cleanly. Lets chains include bias /
scale / eps / mask factors that previously broke the chain at
a Binary op with mismatched shapes.
Fields
prologue: RegionPrologueFKL-style closed fusion: apply before the element-wise chain.
TransformRegion
Fused transform chain (resize, future crop/color). Decompose via
rlx_fusion::DecomposeFusionRegions when no native kernel exists.
BatchElementwiseRegion
Identical Op::ElementwiseRegion chain over num_batch_inputs tensors
(horizontal / z-plane fusion). Inputs are separate batch slices.
Fields
prologue: RegionPrologueMatMul
Matrix multiply. Inputs: [.., M, K] × [.., K, N] → [.., M, N]. Batch dimensions are broadcast.
DotGeneral
Matrix multiply with explicit dimension specification. Like XLA’s DotGeneral — handles arbitrary batch/contracting dims.
Fields
BatchedDenseSolve
Batched dense linear solve. Inputs: A [B, N, N],
b [B, N] or b [B, N, K]. Output: same shape as b.
Per-batch independent solve — each A[i] and b[i] are
solved as a separate Op::DenseSolve. Emitted by vmap of
Op::DenseSolve. The CPU lowering loops over the batch
dimension calling dgesv per slice (LAPACK doesn’t expose a
batched solve on Accelerate; cuSOLVER does on NVIDIA).
DenseSolve
Dense linear solve x = A⁻¹ · b via LU factorization.
Inputs: A [N, N], b [N] (or b [N, K] for multi-RHS).
Output: same shape as b.
VJP via the implicit-function theorem:
dx = solve(Aᵀ, upstream)
dA = -outer(dx, x) (x is the forward output)
db = dx
The rule is dtype-agnostic; lowering is per-backend (Accelerate
dgesv / sgesv, cuSOLVER, etc.).
Cholesky
Cholesky factorization of a symmetric positive-definite matrix.
Input A is [n, n] (SPD, row-major); output L is [n, n]
lower-triangular with the strict upper triangle zeroed, so A = L·Lᵀ.
One input. CPU: LAPACK potrf (promoted to f64). Backward (added with
TriangularSolve) is Ā = Φ(L⁻ᵀ · Φ(Lᵀ L̄) · L⁻¹) symmetrized.
TriangularSolve
Solve op(A)·X = B for triangular A ([n, n]), with B [n] or
[n, nrhs]. op(A) is Aᵀ when transpose else A; A is read as
lower- or upper-triangular per lower (non-unit diagonal). Two inputs
[A, B]; output X matches B. CPU: BLAS trsm. Differentiable in
both (grads are themselves triangular solves).
Det
Determinant of a square matrix A [n, n] → scalar []. One input.
CPU: LAPACK LU (getrf). Backward: Ā = (det(A)·ū)·A⁻ᵀ.
LogDet
log|det(A)| of a square matrix A [n, n] → scalar []. One input.
Numerically stable (sum of log|Uᵢᵢ|). Backward: Ā = ū·A⁻ᵀ.
Sort
Sort values along axis (ascending, or descending). One input; output
same shape/dtype. Backward scatters ū back via the argsort permutation.
ArgSort
Indices that would sort along axis, as f32 (like ArgMax). One input;
output same shape. Non-differentiable.
Svd
One factor of a thin SVD A = U·diag(S)·Vᵀ (A [m, n],
k = min(m, n)), selected by part. One input. CPU: LAPACK gesdd.
Only SvdPart::S is differentiable (Ā = U·diag(ū)·Vᵀ); U/Vt are
forward-only.
Qr
One factor of a thin QR A = Q·R (A [m, n], k = min(m, n)),
selected by part: Q [m, k] or R [k, n]. One input. CPU: LAPACK
geqrf+orgqr. Differentiable (the QR pullback via copyltu + R⁻ᵀ).
LayerNorm
Layer normalization: input, gamma, beta → normalized output.
axis is the feature dimension (usually -1).
GroupNorm
Group normalization on NCHW tensors: input, gamma, beta → same shape.
Normalizes over (C/num_groups) × H × W per group.
LayerNorm2d
LayerNorm2d on NCHW: normalize across the channel axis at each spatial
position (candle / SAM LayerNorm2d semantics — not PyTorch’s H×W norm).
ResizeNearest2x
Nearest-neighbor 2× upsample on NCHW (doubles spatial dims 2 and 3).
Interpolate3d
Nearest-neighbor resample on NCDHW. size is [D_out, H_out, W_out].
Source mapping: src = min(floor(dst * in / out), in - 1).
RmsNorm
RMS normalization: input, gamma → normalized output.
BatchNormInference
BatchNorm inference with frozen running statistics.
Inputs: x, gamma, beta, running_mean, running_var.
Feature dimension is the last axis of x; stats are 1-D [C].
Attention
Scaled dot-product attention: Q, K, V, [mask] → output.
The compiler can lower this to fused SDPA or flash attention.
mask_kind controls how masking is applied — Custom reads from
the 4th input tensor; None / Causal / SlidingWindow skip the
mask load and apply the mask directly in the inner loop. See
MaskKind for the rationale.
score_scale: when Some(s), dot-product scores are multiplied by
s instead of the default 1/sqrt(head_dim) (Gemma uses head_dim^-0.5
explicitly in config). attn_logit_softcap: when Some(c), applies
c * tanh(s/c) to scores before softmax (Gemma 2).
§Operand layout
Q/K/V (and the output) carry the heads explicitly and are accepted in two
rank-4 layouts — both end in head_dim, so backends disambiguate by which
axis equals num_heads:
[B, S, H, D](heads at axis 2) — the dominant convention, produced by Llama/Moshi-stylereshape([B, S, H, D])after the QKV projection. The CPU/Metal/MLX/wgpu kernels attend over axis 1 (S); CoreML transposes to[B, H, S, D]first (seerlx-coreml).[B, H, S, D](heads at axis 1) — already head-canonical.
Rank-3 [B, S, D] is single-head and head-canonical. K/V may have a
different sequence length than Q (S_k >= S_q) for KV-cache decode; the
causal mask places the S_q queries at the end of the S_k keys
(query qi is at absolute position S_k - S_q + qi).
Fields
v_head_dim: Option<usize>Value/output per-head width. None ⇒ same as head_dim (the common
case). Some(v) with v != head_dim is asymmetric SDPA — Q/K use
head_dim for the scores, V is read with v_head_dim and the output
is num_heads * v_head_dim wide. Used by DeepSeek/Kimi MLA, which
otherwise zero-pads V to head_dim and slices the output back.
Rope
Rotary position embedding applied to one tensor: x, cos, sin → x_rotated.
Apply separately to Q and K. head_dim is the per-head width; n_rot
is how many leading dims get NeoX RoPE (pair offset n_rot/2). When
n_rot < head_dim, trailing dims are copied unchanged (Qwen3.5 MRoPE).
Fields
style: RopeStyleRotation pairing convention (default RopeStyle::NeoX).
AxialRope2d
SAM2 axial 2-D RoPE on [batch, seq, num_heads * head_dim].
Reshape
Transpose
Narrow
Select a contiguous slice along an axis.
Slice
Strided slice along an axis: output element j reads input index
start + j*step for j in 0..len. step may be negative (reverse);
step != 0. All accessed indices must lie in [0, axis_len). The
contiguous step == 1 case is Op::Narrow; this covers x[::2],
x[::-1], dilation, and general strided indexing. No backend has a
native kernel except Metal/CUDA; others legalize to narrow/reverse/
gather via rlx_fusion::LowerSlice.
Roll
Cyclic shift along one or more axes — jnp.roll / torch.roll.
out[.., i, ..] = x[.., (i - shift) mod n, ..] for each (shift, dim)
pair, applied in order. Shifts may be negative or exceed the axis length;
both are reduced modulo n. Output shape always equals input shape —
nothing is dropped or introduced, which is what separates this from
Op::Slice and from a circular Op::Pad.
No backend has a native kernel: rlx_fusion::LowerRoll legalizes each
axis to concat(narrow(tail), narrow(head)), which every backend already
supports. A native kernel is worth adding only if a workload shows the
two-copy decomposition is hot — the copies are the whole cost, so a fused
version saves one pass, not an order of magnitude.
Fields
Concat
Concatenate along an axis.
KvAppend
In-place KV-cache append: write row (input 1) into cache (input 0)
at index pos along axis, and return the [..pos+1..] prefix of
cache — the output aliases cache’s buffer (no copy). Replaces the
decode concat(past_kv, new_row) whose full-cache copy is O(context);
this writes a single row (O(1)). The planner assigns the output cache’s
slot (see collect_view_aliases); backends encode a 1-row write.
Expand
Expand (broadcast) to a target shape.
Gather
Gather elements by index along an axis (embedding lookup).
Reverse
Reverse (flip) the element order along each listed axis. Output shape
equals the input shape; element […, i, …] on a reversed axis of size
d reads input […, d-1-i, …]. Batch-general — only the listed axes
flip, all others pass through unchanged (the right primitive for
reversing a [batch, seq, …] sequence without a batch=1 assumption).
Pad
Pad each axis by pads[i] = [before, after] (aligned to the input rank)
per PadMode. Output axis i grows by before + after. Single tensor
input; the constant fill value lives in the mode, not an input.
No backend has a native kernel except Metal/CUDA; every other backend
legalizes it to narrow/reverse/concat/gather/full via
rlx_fusion::LowerPad, so the decomposition is the semantic oracle.
Clamp
Elementwise clamp(x, min, max). Decomposes to max/min (native on
every backend); no dedicated kernel needed. Backward: pass-through where
min ≤ x ≤ max, else 0.
Tile
Tile (repeat) the input reps[i] times along axis i (reps aligned to
the input rank). Decomposes to concat. Backward: reduce-sum the copies.
Trilu
Keep the upper (upper=true) or lower triangle relative to diagonal;
zero the rest. Operates on the last two axes. Decomposes to a multiply by
a constant triangular mask. Backward: same mask.
Reduce
Reduce along specified axes.
Histogram
Histogram of x into bins equal-width buckets over [min, max].
Output is a 1-D f32 tensor [bins] of counts. Elements outside
[min, max] are dropped; x == max lands in the last bin (half-open
buckets, closed at the top — matches numpy.histogram). Native on CPU;
decomposes to Compare+Reduce::Sum on every other backend (see
LowerHistogram). Non-differentiable (discrete counts, zero gradient).
Introduced for the data-pattern recording harness.
SelectiveScan
Selective scan (plan #15) — Mamba-style state-space model
step. The recurrence:
h[t] = exp(Δ[t] * A) * h[t-1] + Δ[t] * B[t] * x[t]
y[t] = C[t] * h[t]
where state h has dimension state_size and the input has
(batch, seq, hidden).
Inputs (in order):
x [b, s, h] f32 input
delta [b, s, h] f32 step size (per-position, per-channel)
a [h, n] f32 transition matrix (one per channel)
b [b, s, n] f32 input projection
c [b, s, n] f32 output projection
Output: [b, s, h] f32. State h is implicit; the kernel
scans through the seq dimension carrying it.
state_size = n is exposed for the cost model.
GatedDeltaNet
Gated DeltaNet linear-attention recurrence — the per-layer
kernel used by Qwen3.5/3.6 trunk “linear attention” blocks
(and Qwen3-Next, Kimi-Linear). Mirrors
llama.cpp / src/models/delta-net-base.cpp autoregressive
path; chunked + fused variants ride the same op identity.
Math (per token t, head h, state size n):
state matrix S[h, i, j] is implicit (reset per batch).
S[h] *= exp(g[t,h]) # scalar gate
sk[h,j] = Σ_i S[h,i,j] * k[t,h,i]
d[h,j] = (v[t,h,j] - sk[h,j]) * b[t,h] # b = beta
S[h,i,j] += k[t,h,i] * d[h,j] # outer-prod
o[t,h,j] = Σ_i S[h,i,j] * (q[t,h,i] / √n)Inputs:
q [b, s, h_v, n] f32 queries (L2-normed by caller)
k [b, s, h_v, n] f32 keys (L2-normed by caller;
GQA-repeated to match h_v)
v [b, s, h_v, n] f32 values
g [b, s, h_v] f32 log-gate (exp’d inside kernel)
beta [b, s, h_v] f32 delta-rule mixing factor
Output: [b, s, h_v, n] f32.
When carry_state is true, a sixth input state [b, h_v, n, n]
provides the initial SSM matrix per head; the kernel updates it
in place across the sequence and leaves the final state in the
same buffer (same layout as the internal scan state:
state[h, i, j] row-major over (n, n) per head).
When gate_per_channel is true (Kimi-K3 “KDA”), the g input is
[b, s, h_v, n] (one log-gate per key channel) instead of [b, s, h_v],
and the state decay is applied per key-row: S[i, j] *= exp(g[i])
(vs the scalar S *= exp(g) of the per-head gate). Everything else is
identical; Qwen3-Next / Qwen3.5 use the per-head gate (false).
GatedDeltaNetBackward
Fused backward of Op::GatedDeltaNet, computing every input gradient
from one reverse scan.
Without it the gradient walk has to unfuse the forward: the time loop is
unrolled into per-timestep primitives, which for a 24-token prompt at
state_size = 128 is ~585 nodes and runs ~32× slower than the fused
kernel. The reason is structural rather than dispatch overhead — in SSA
graph form every timestep materializes a fresh [B·H, N, N] state,
while the kernel updates one working set in place.
Inputs are [q, k, v, g, beta], + [state] when carry_state, then
dy last — so the optional state keeps index 5, matching the forward.
The gradients have different shapes and a node has one output, so they
come back packed into a 1-D tensor, in this order: dq, dk, dv
(each B·S·H·N), dg (B·S·H·N when gate_per_channel, else
B·S·H), dbeta (B·S·H), and dstate (B·H·N·N, only when
carry_state). Op::Narrow + Op::Reshape recover each one; the
VJP rule emits exactly that.
Lstm
Multi-layer (optionally bidirectional) LSTM over a
[batch, seq, input] sequence. Gate order i, f, g, o (PyTorch);
recurrence per step:
z = x_t · w_ihᵀ + h_{t-1} · w_hhᵀ + bias
i,f,o = σ(z_i), σ(z_f), σ(z_o); g = tanh(z_g)
c_t = f · c_{t-1} + i · g; h_t = o · tanh(c_t)D = 2 if bidirectional else 1. Inputs [x, w_ih, w_hh, bias]
(+ [h0, c0] when carry):
x:[batch, seq, input]w_ih: packed, all(layer, direction)blocks concatenated inlayer-major thendirectionorder. Block(l,d)is[4*hidden, in_l]within_l = inputforl=0elseD*hidden.w_hh: packedL*Dblocks of[4*hidden, hidden].bias: packedL*Dblocks of[4*hidden](combinedb_ih+b_hh).h0,c0(carry only):[L*D, batch, hidden]; the finalhn/cnare written back in place (decode threading).
Output y: [batch, seq, D*hidden] — last layer’s hidden states
(forward ‖ reverse concatenated on the feature axis). With
num_layers = 1, bidirectional = false, carry = false this is the
plain single-layer LSTM and the weight shapes reduce to
[4*hidden, input], [4*hidden, hidden], [4*hidden].
Gru
Gated Recurrent Unit (PyTorch). D = 2 if bidirectional else 1;
gate order r, z, n. Recurrence:
r = σ(x·W_irᵀ + b_ir + h·W_hrᵀ + b_hr)
z = σ(x·W_izᵀ + b_iz + h·W_hzᵀ + b_hz)
n = tanh(x·W_inᵀ + b_in + r ⊙ (h·W_hnᵀ + b_hn))
h' = (1 - z) ⊙ n + z ⊙ hThe new-gate reset is applied to the hidden term after its bias,
so b_ih/b_hh cannot be merged. Inputs [x, w_ih, w_hh, b_ih, b_hh]
(+ [h0] when carry); packing matches Op::Lstm with 3*hidden
gate rows. h0 [L*D, batch, hidden]. Output [batch, seq, D*hidden].
Rnn
Elman RNN (PyTorch): h' = act(x·w_ihᵀ + h·w_hhᵀ + bias) with
act = ReLU when relu else tanh. Packed w_ih [hidden, in_l],
w_hh [hidden, hidden], merged bias [hidden] (per layer ×
direction). Inputs [x, w_ih, w_hh, bias] (+ [h0] when carry).
Output [batch, seq, D*hidden].
Mamba2
Mamba-2 / SSD (structured state-space duality) scan — the
scalar-decay SSM at the core of Mamba-2, sibling of
Op::SelectiveScan / Op::GatedDeltaNet. Inputs
[x, dt, a, b, c] (all f32):
x:[batch, seq, heads, head_dim]dt:[batch, seq, heads]— discretization step (already ≥ 0, e.g. softplus’d by the caller)a:[heads]— per-head decay rate (dA = exp(dt·a))b:[batch, seq, heads, state_size]c:[batch, seq, heads, state_size]
State S [batch, heads, head_dim, state_size] is zero-initialized
per sequence. Recurrence per timestep t:
dA_t = exp(dt_t · a)
S_t = dA_t · S_{t-1} + (dt_t · x_t) ⊗ b_t
y_t = Σ_n S_t[:, n] · c_t[n]Output y: [batch, seq, heads, head_dim] (same shape as x).
DequantMatMul
Fused dequant + matmul (plan #5). The biggest LLM-bandwidth win on Apple Silicon: dequantizes weights inside the matmul inner loop, never materializing f32 weights.
BREAKING CHANGE in 0.2.0: num_inputs() is now
scheme-dependent — 4 for legacy Int8 schemes, 2 for
the new GGUF K-quant schemes (their scales/mins live inside
the packed bytes, so no side-channel scale / zp tensors
are fed in). Callers that assumed a fixed 4-input contract
must inspect scheme.is_gguf() before reading inputs.
Inputs (Int8 schemes — scheme.is_gguf() == false):
x [m, k] f32 activations
w_q [k, n] packed quantized weight bytes (i8 per
element for Int8 schemes; 4-bit
packed two-per-byte for Int4)
scale [k/block, n] per-block f32 dequant scale
zp [k/block, n] per-block f32 zero-point
(zero-tensor if symmetric)
Inputs (Nvfp4Block — fixed group size 16 along K):
x [m, k] f32 activations
w_q [k,n/2] packed FP4 E2M1 codes (unsigned nibble 0..15)
scale [k/16, n] u8 FP8 E4M3 block scales (one byte / group)
global_scale [1] f32 per-tensor scale (pass [1.0] if unused)
Inputs (GGUF schemes — scheme.is_gguf() == true):
x [m, k] f32 activations
packed_w [bytes] raw GGUF super-block bytes; the
dequantizer reads the per-sub-block
scales / mins / quants directly out
of the buffer per the K-quant block
layout (no side tensors).
Output: [m, n] f32.
block_size (on the Int8 schemes only) is the number of
consecutive elements that share one (scale, zero_point) pair.
The Op carries enough metadata that the kernel doesn’t need
a separate QuantMap lookup at run time.
Fields
scheme: QuantSchemeSynthMatMul
On-chip weight synthesis matmul: reconstruct the weight from a
compact learned representation inside the matmul inner loop,
trading DRAM weight-bandwidth for ALU — the sibling of
Op::DequantMatMul, but with a learned function rather than a
fixed quant grid. Inputs (in order), for SynthKind::Codebook:
x[m, k]f32 activationsindices[n, k/entry_dim]u8 codebook indices (weight stored transposed)codebook[num_entries, entry_dim]f32 learned centroids Outputy:[m, n]f32,y = x · WᵀwhereW[n, kb·entry_dim + t] = codebook[indices[n, kb]][t].
SynthMatMulBackward
Fused backward of Op::SynthMatMul — one dispatch per gradient instead
of the ~11 primitives the generic VJP would emit (Gather + 2 MatMul +
Transpose + ScatterAdd + reshapes). Inputs (in order):
x[m, k]f32 activations (forward input)indices[n, k/entry_dim]u8 codebook indicescodebook[num_entries, entry_dim]f32 centroidsupstream[m, n]f32 grad w.r.t. the forward output Output isdx [m,k]forwrt=Dx, ord_codebook [num_entries,entry_dim]forwrt=Codebook. Backends without a native kernel decompose it viaLowerSynthMatMulBackward(bit-identical to the old primitive expansion).
SynthReconstruct
Reconstruct the dense weight W [k,n] from indices [n, k/entry_dim] (u8)
codebook [num_entries, entry_dim], in ONE fused kernel (no separate cast/gather/reshape/transpose). Lets a model fold multi-stage VQ + LoRA into oneW_eff = Σ reconstruct + A·Bᵀthen a single matmul. Output[k,n]:W[kb·entry_dim+t, j] = codebook[indices[j, kb], t]. Backends without the native kernel decompose viaLowerSynthReconstruct.
SplineActivation
KAN (Kolmogorov–Arnold) learnable spline activation: each channel c
gets its own univariate function φ_c expressed in a fixed Gaussian-RBF
basis, with learned per-channel coefficients — “weights as functions”.
Inputs:
x[.., C]f32 pre-activationscoeff[C, num_basis]f32 learned basis coefficients Outputy: same shape asx, withy[.., c] = Σ_g coeff[c, g] · exp(-((x[.., c] − center_g) · inv_h)²), wherecenter_g = grid_min + g·(grid_max−grid_min)/(num_basis−1)andinv_h = (num_basis−1)/(grid_max−grid_min)(adjacent RBFs ~one width apart). Differentiable wrt bothxandcoeff.
SplineActivationBackwardX
Fused KAN spline VJP → dx (input grad). Inputs: x, coeff, upstream.
Builds the RBF basis + derivative in registers (loop over num_basis) and
emits dx = upstream·Σ_g coeff[c,g]·B_g'(x) in one pass — replaces the
decomposed exp/mul/reduce chain that materializes the [.., C, G] basis.
SplineActivationBackwardCoeff
Fused KAN spline VJP → dcoeff (coeff grad [C, num_basis]). Inputs: x,
upstream. Atomic-accumulates dcoeff[c,g] = Σ_batch upstream·B_g(x).
QMatMul
Real INT8-arithmetic matrix multiply with i32 accumulation.
Inputs (in order):
x [M, K] i8 activations (zero-point = x_zp)
w [K, N] i8 weights (zero-point = w_zp)
bias [N] i32 (in accumulator scale = x_scale·w_scale),
pass a zeros tensor for “no bias”
Output: [M, N] i8 (zero-point = out_zp)
Per-element compute:
out[m,n] = requantize(bias[n] + Σₖ (x[m,k]-x_zp)·(w[k,n]-w_zp), mult, out_zp)
where mult = x_scale · w_scale / out_scale.
This is the same kernel shape rlx-cortexm/src/dense.rs
uses for on-device int8 inference, lifted into the IR so the
rlx-cpu backend can run a quantized graph directly (instead
of round-tripping through fake-quant Dequantize → MatMul →
Quantize). 2-D only — generalizing to batched comes when a
real workload demands it.
QConv2d
Real INT8-arithmetic 2-D convolution with i32 accumulation.
Inputs:
x [N, C_in, H, W] i8 (zero-point = x_zp)
w [C_out, C_in/groups, kH, kW] i8 (zero-point = w_zp)
bias [C_out] i32 in accumulator scale
Output: [N, C_out, H_out, W_out] i8 (zero-point = out_zp).
Same NCHW geometry contract as Op::Conv; same requantize
math as Op::QMatMul (per-element acc·mult rounded to i8).
Fields
ScaledMatMul
Native low-precision tensor-core GEMM with f32 accumulation —
both operands are sub-f16 codes fed directly into the matmul, the
opposite of DequantMatMul (which decodes weights to f32 first and
keeps x in f32). This is the path that wins on Hopper / Ada /
Blackwell (cublasLt FP8/FP4) and CDNA3 / CDNA4 (hipBLASLt). On CPU /
Metal it is the decode-and-accumulate reference (no fp8 matrix HW).
Layout is TN — both operands are K-last (out = lhs · rhsᵀ) so
every block scale runs along the last/contraction axis uniformly, which
is also cuBLASLt / hipBLASLt FP8’s preferred layout.
Inputs (in order):
lhs [m, k] U8 — packed crate::ScaledFormat codes (logical dims)
rhs [n, k] U8 — packed codes (K-last, i.e. weight transposed)
lhs_s scale tensor, dtype per crate::ScaleLayout::scale_dtype
rhs_s scale tensor
bias [n] f32, present iff has_bias
Output: [m, n] f32. Reconstruction: decode(code)·scale(block).
Operand Shapes carry logical element dims with dtype = U8; the
physical byte layout (1 code/byte on the CPU oracle, packed nibbles on
GPU) is a backend-paired detail shared by ScaledQuantize and the
matmul kernel, so it never needs a DType variant.
ScaledQuantize
Quantize an f32/f16 tensor to packed low-precision codes for
Op::ScaledMatMul. Inputs: x (f32) and scale (from
Op::ScaledQuantScale); output is DType::U8 codes of the same
logical shape as x. code[i] = encode(format, x[i] / scale(block_of(i))).
ScaledQuantScale
Compute the (per-tensor or per-block) scale tensor for a tensor about
to be quantized to format. Input: x (f32). Output dtype follows
crate::ScaleLayout::scale_dtype. Split from ScaledQuantize to keep
the IR single-output (mirrors the FakeQuantizeLSQBackwardX/Scale split).
ScaledDequantize
Reconstruct f32 from packed low-precision codes: decode(code)·scale.
The inverse of Op::ScaledQuantize. Inputs: codes (U8), scale;
output is f32 with the same logical shape as codes. Used by the
straight-through VJP of Op::ScaledMatMul to rebuild operands in the
backward graph (and as a standalone dequantizer).
LoraMatMul
Fused LoRA matmul: out = x·W + scale * x·A·B.
Inputs (in order): x [m, k], w [k, n], a [k, r], b [r, n].
r is the LoRA rank (typically 4-64). scale is the
per-adapter alpha / rank knob.
Plan #9: lifts LoRA from “three matmuls + an add” into one
kernel that keeps the rank-r intermediate in registers.
PartitionedConv
Uniform-partitioned overlap-save 1-D convolution (FIR / RIR).
Inputs [x, ir]: signal x [.., L] and rank-1 impulse response
ir [M]; output is the full linear convolution [.., L + M − 1].
block sets the partition / FFT size (2·next_pow2(block) points).
Decomposes (via unfuse) to rfft → batched complex matmul over the partition axis → irfft, so the frequency-domain delay line runs on the
native batched-GEMM kernels (cuBLAS / rocBLAS / MPS). See
crate::Graph::partitioned_conv.
Sample
Fused sampling kernel: logits → optional top-k filter →
optional top-p truncation → softmax → multinomial sample.
One f32-encoded sampled token id per batch row (output
shape [batch]).
temperature == 1.0 matches a plain argmax-of-softmax;
lower → sharper, higher → flatter. top_k == 0 disables.
top_p == 1.0 disables. seed is the Philox seed; pass 0
for “use process-global counter” (still deterministic
given the call order).
Borrowed from MAX’s nn/sampling.mojo.
Latency-critical: never materializes the full softmax
distribution on the host.
RngNormal
ONNX RandomNormalLike / RandomNormal: zero or one shape-template
input (Like uses the template’s shape; RandomNormal with a shape
attribute needs no input). Output shape is fixed on the node.
at compile/execute time. Optional ONNX seed attribute (f32) overrides
the mixed seed on the Ort backend.
RngUniform
ONNX RandomUniformLike.
Cumsum
Inclusive cumulative sum along an axis. Same shape in/out.
Underpins ragged-tensor offsets, sampling (top-p prefix sum),
and sequence-position math.
exclusive=true shifts the result so output[0] = 0 (useful
for offset arrays where the first segment starts at 0).
CumProd
Inclusive cumulative product along an axis. Same shape in/out.
exclusive=true shifts so output[0] = 1 (the multiplicative identity).
Mirrors Op::Cumsum; native on CPU/Metal/CUDA, decomposes to a
masked reduce-product elsewhere.
CumMax
Inclusive cumulative maximum along an axis. Same shape in/out.
exclusive=true shifts so output[0] = -inf. Native on CPU/Metal/CUDA,
decomposes to a masked reduce-max elsewhere.
ArgMax
Index of the maximum along axis. Output drops that axis (or keeps it
as size 1 when keep_dim). Indices are f32-encoded (rlx is f32 at
the I/O boundary, matching Op::TopK).
ArgMin
Index of the minimum along axis; see Op::ArgMax.
Softmax
Softmax along an axis (reduction + element-wise).
TopK
Top-K indices along the last axis. Output shape [..., k],
f32-encoded indices (rlx is f32-only at the I/O boundary).
To recover the values, follow with a Gather against the
original tensor — works because Gather already supports any axis.
Ties broken by smaller index (matches NumPy / PyTorch
torch.topk(..., largest=True, sorted=True)).
Used by MoE gating; also useful for beam search.
GroupedMatMul
Indexed batched matmul. The MoE GEMM primitive.
Inputs: [input, weight, expert_idx]
input : [M, K] — per-token activations
weight : [num_experts, K, N] — stacked expert weights
expert_idx : [M] — f32-encoded expert id per token
Output : [M, N] — output[i] = input[i] @ weight[expert_idx[i]]
CPU is a real segmented GEMM: counting-sort tokens by expert, one GEMM per
expert, then unpermute. GPU backends dispatch a dedicated grouped kernel.
DequantGroupedMatMul
Fused GGUF K-quant dequant + Op::GroupedMatMul. Same three
inputs as GroupedMatMul, but weight is a U8 tensor holding
num_experts contiguous packed expert slabs (GGML layout, expert
dimension outermost). Scales live inside the packed bytes.
Fields
scheme: QuantSchemeDequantGroupedMatMulMlx
MLX-affine grouped expert matmul — the Op::DequantGroupedMatMul
analogue for mlx-community MoE, whose experts carry SEPARATE scales
and biases tensors (not embedded in the packed bytes). Five inputs:
input : [M, K] — one row per token
weight : [E · out · (K·bits/32)] U8 — stacked affine code words
scales : [E, out, n_groups] f32 — per-expert per-group scales
biases : [E, out, n_groups] f32 — per-expert per-group biases
expert_idx : [M] — f32-encoded expert id per row
Output : [M, out] — output[i] = input[i] @ dequant(expert_idx[i])
Fields
scheme: QuantSchemeDequantMoEWeights
Dequant a packed MoE expert stack to F32 [num_experts, K, N] in
GroupedMatMul layout. Input: U8 packed bytes; output shape is
declared on the node ([E, K, N]).
Fields
scheme: QuantSchemeScaledGroupedMatMul
Native low-precision grouped (MoE) tensor-core GEMM — the
expert-indexed analogue of Op::ScaledMatMul. Both operands are
sub-f16 crate::ScaledFormat codes fed directly into a segmented
GEMM with per-block (scale_layout) rescaling, so the MXFP4 path
(F4E2M1 + BlockMxE8M0) wins on Blackwell / CDNA4 (native FP4
tensor cores) while every other backend runs the decode-and-segment
reference (see the CPU oracle and LowerScaledGroupedMatMul).
Layout is TN per expert (out = input · weight[e]ᵀ) so every
block scale runs along the K/contraction axis, matching ScaledMatMul.
Inputs (in order):
input [M, K] U8 — packed codes (per-token acts)
weight [E, N, K] U8 — packed codes, one slab per expert (K-last)
input_s [M, ⌈K/block⌉] scale tensor, dtype per crate::ScaleLayout::scale_dtype
weight_s [E, N, ⌈K/block⌉] per-expert scale tensor
expert_idx [M] f32-encoded expert id per token
bias [E, N] f32, present iff has_bias (per-expert)
Output: [M, N] f32. out[i] = decode(input[i])·decode(weight[eidx[i]])ᵀ.
Output shape [M, N] is derived from input/weight in infer_shape.
ScatterAdd
Scatter-add along axis. The “unpermute” half of MoE routing, the
transpose of Op::Gather, and the shape every immersed-boundary or
segment-reduction kernel wants (many grid cells accumulating into few
slots).
Inputs: [updates, indices]
updates : same rank as the output, with axis of length num_updates
indices : [num_updates] — destination position along axis
Output : node’s declared shape; out[.., indices[i], ..] += updates[.., i, ..]
Initial output is zero; repeated indices accumulate (sequentially on CPU, atomic-add on GPU).
Indices must lie in [0, out_dim) — out of range is a caller error, not
a supported “drop these” mode. The CPU kernel debug_asserts it; other
backends are unchecked. This contract predates the axis parameter and
is unchanged by it.
§Index dtype
indices may be crate::DType::I64 (matching Op::Gather) or
F32. The F32 encoding is historical and exact only below 2^24;
crate::verify rejects it when the destination is larger, since past
that point distinct rows share a representation and updates land on the
wrong slot with no error.
§Axis support
Backends implement axis == 0 only. rlx_fusion::LowerScatterAddAxis
rewrites any other axis to transpose → scatter → transpose and
normalizes I64 indices to the F32 the kernels read, so a backend
never sees a form it does not handle. Adding a native strided kernel is
a per-backend optimization, not a correctness requirement.
ScatterNd
ONNX ScatterND: copy data, then write updates at multi-index
locations from indices.
Inputs: [data, indices, updates]
data : rank r tensor (output shape matches)
indices : rank q, trailing dim k indexes the first k axes of data
updates : shape indices.shape[:-1] + data.shape[k:]
reduction selects overwrite vs accumulate (see ScatterNdReduction).
Fields
reduction: ScatterNdReductionScatterElements
ONNX ScatterElements: per-element scatter along axis.
Inputs: [data, indices, updates] — output shape matches data.
indices and updates share rank with data; along axis they
select which coordinate to write.
GatherNd
ONNX GatherND: gather slices from data indexed by the trailing
axis of indices. Inputs: [data, indices].
batch_dims (default 0) skips leading shared batch axes.
GatherElements
ONNX GatherElements / take_along_axis: output shape = indices
shape; each output element is data with the axis coordinate
replaced by the corresponding index. Inputs: [data, indices].
Conv
2D convolution on NCHW tensors. Also exposed as OpKind::Conv / conv2d.
Weight layout: [C_out, C_in / groups, kH, kW].
Fields
Im2Col
NCHW im2col for conv backward-weight style matmul.
Input [N, C, H, W]. Output [M, C·kH·kW] with
M = N · H_out · W_out. When batch is dynamic::sym::BATCH,
output rows use dynamic::sym::ROWS (bind N · H_out · W_out).
ConvTranspose2d
2D transposed convolution on NCHW. Weight layout (PyTorch):
[C_in, C_out / groups, kH, kW].
Fields
Conv3d
3D convolution on NCDHW tensors (forward / cross-correlation).
Weight layout: [C_out, C_in / groups, kD, kH, kW]. Kernel size is
derived from the weight shape. Mirrors Op::Conv with a depth axis.
ConvTranspose3d
3D transposed convolution on NCDHW. Weight layout (PyTorch):
[C_in, C_out / groups, kD, kH, kW]. Mirrors Op::ConvTranspose2d.
Fields
Pool
ReluBackward
ReLU backward: dx = dy where x > 0 else 0.
Inputs: [x, dy] — both same shape; output matches.
ComplexNormSq
Element-wise complex squared-magnitude: |z|² = z.re² + z.im².
Input: 1 tensor with DType::C64. Output: same shape but
DType::F32. The natural real-valued loss surface for
Wirtinger reverse-mode AD on complex graphs — pair with
Op::ComplexNormSqBackward.
Conjugate
Element-wise complex conjugate: z̄ = z.re - i·z.im per element.
Input: 1 tensor with DType::C64. Output: same shape, same dtype.
Used by Wirtinger VJP rules on Op::Binary over C64 (the rule
for y = a·b is dL/dā = upstream · conj(b), etc.).
ComplexNormSqBackward
Backward for Op::ComplexNormSq under Wirtinger calculus.
f(z) = |z|² = z·z̄, so ∂f/∂z̄ = z. Given upstream real
cotangent g (same shape as the forward output), the C64
gradient with respect to z is g · z element-wise, returned
in C64 storage [re_g·re_z, re_g·im_z] per element.
Inputs: [z (C64), g (F32)] — both same logical shape; output
matches z (C64).
LayerNormBackwardInput
LayerNorm backward w.r.t. the input. Computes
d_x[..., d] = inv_std · (dy·γ - mean(dy·γ) - x̂·mean(dy·γ·x̂))
over the feature axis, where x̂ = (x - mean)/std is recomputed
inline from x. Inputs: [x, gamma, dy]; output shape = x.shape.
Currently lowers axis=-1 only (matches the forward thunk).
LayerNormBackwardGamma
LayerNorm backward w.r.t. gamma. Computes
d_gamma[d] = Σ_{batch} dy[..., d] · x̂[..., d]
— sums the per-element product of upstream and the (recomputed)
normalized input over the leading axes. Inputs: [x, dy];
output shape = gamma.shape (= 1-D feature axis).
RmsNormBackwardInput
RMSNorm backward w.r.t. input. Inputs [x, gamma, beta, dy]; output = x.shape.
RmsNormBackwardGamma
RMSNorm backward w.r.t. gamma. Inputs [x, gamma, beta, dy]; output = gamma.shape.
RmsNormBackwardBeta
RMSNorm backward w.r.t. beta. Inputs [x, gamma, beta, dy]; output = beta.shape.
RopeBackward
RoPE backward w.r.t. x. Inputs [dy, cos, sin]; output = dy.shape.
Fields
style: RopeStyleRotation pairing convention — must MATCH the forward Op::Rope
this is the adjoint of.
This field is load-bearing and used not to exist. vjp_rope
destructured the forward as Op::Rope { head_dim, n_rot, .. }, and
that .. silently dropped style: every RopeStyle::GptJ rotation
— the GGUF convention — got a NeoX adjoint. The forward was correct,
so nothing downstream complained; the gradient was simply wrong, on
every backend at once, for full and partial rotation alike.
rlx-runtime’s fd_backward_gate pins both styles now.
GroupNormBackwardInput
GroupNorm (NCHW) backward w.r.t. input. Inputs [x, gamma, beta, dy].
GroupNormBackwardGamma
GroupNorm backward w.r.t. gamma. Inputs [x, dy]; output = gamma.shape.
GroupNormBackwardBeta
GroupNorm backward w.r.t. beta. Inputs [x, dy]; output = beta.shape.
BatchNormInferenceBackwardInput
BatchNorm inference backward w.r.t. x. Inputs [x, gamma, mean, var, dy].
BatchNormInferenceBackwardGamma
BatchNorm inference backward w.r.t. gamma. Inputs [x, mean, var, dy].
BatchNormInferenceBackwardBeta
BatchNorm inference backward w.r.t. beta. Inputs [dy]; output = beta.shape.
CumsumBackward
Cumsum backward along axis. Inputs [dy]; output matches forward input shape.
GatherBackward
Gather backward (scatter-add into table). Inputs [dy, indices]; output = table shape.
axis matches forward Op::Gather.
ActivationBackward
Generic element-wise activation backward. kind selects the
closed-form derivative d/dx act(x). Inputs: [x, dy]; output
shape matches x. The kernel computes d/dx · dy per element.
Closed forms (all element-wise):
Gelu— exact derivative of erf-based GELU.GeluApprox— derivative of the tanh approximation0.5 x (1 + tanh(√(2/π)(x + 0.044715 x³))).Silu—σ(x)·(1 + x·(1 - σ(x))).Sigmoid—σ(x)·(1 - σ(x)).Tanh—1 - tanh(x)².Exp—exp(x).Log—1 / x.Sqrt—0.5 / sqrt(x).Rsqrt—-0.5 · x^(-3/2).Neg—-1.Abs—sign(x)(zero at x=0).Sin—cos(x).Cos—-sin(x).Tan—1 + tan²(x).Atan—1 / (1 + x²).Relu— kept here for completeness; the dedicatedReluBackwardop is preferred for relu and is what the autodiff pass actually emits.
Fields
kind: ActivationFakeQuantizeBackward
Backward for Op::FakeQuantize under a non-default STE.
Inputs [x, dy]: the forward input and the upstream
gradient. Output dx same shape. The bits/axis/ste
fields must match the forward op so the kernel computes the
same per-channel scale and applies the right STE attenuation.
For SteKind::Identity this op is unnecessary — autodiff
just routes upstream through unchanged.
MaxPool2dBackward
2D max-pool backward. Routes each element of dy back into the
position in x’s window where the forward max was taken.
Inputs: [x, dy] with x [N, C, H, W] and
dy [N, C, H_out, W_out]. Output: same shape as x.
Carries the forward pool’s geometry so the kernel can recompute
the argmax position per window without a saved-indices tensor.
Conv2dBackwardInput
2D conv backward w.r.t. input. Computes dx = conv_transpose(dy, w).
Inputs: [dy, w] with dy [N, C_out, H_out, W_out] and
w [C_out, C_in/groups, kH, kW]. Output: [N, C_in, H, W]
(declared on the node — caller knows the original input shape).
Geometry is the forward conv’s parameters, not the transposed
conv’s.
Fields
Conv2dBackwardWeight
2D conv backward w.r.t. weight. Computes
dw[c_out, c_in, kh, kw] = sum_{n,h_out,w_out} x[n,c_in,...] * dy[n,c_out,h_out,w_out].
Inputs: [x, dy]. Output: [C_out, C_in/groups, kH, kW].
Fields
MaxPool3dBackward
3D max-pool backward on NCDHW. Inputs: [x, dy]. Output: same as x.
Conv3dBackwardInput
3D conv backward w.r.t. input (NCDHW). Inputs: [dy, w].
Fields
Conv3dBackwardWeight
3D conv backward w.r.t. weight (NCDHW). Inputs: [x, dy].
Fields
SoftmaxCrossEntropy
Fused softmax + cross-entropy against a dense target
distribution (soft labels / one-hot probabilities) — the
companion to the sparse integer-label
Op::SoftmaxCrossEntropyWithLogits. Per-row output:
loss[n] = logsumexp(logits[n]) - Σ_c targets[n,c]·logits[n,c]
= -Σ_c targetsn,c·log_softmax(logitsn)c``. Inputs: [logits, targets], both [N, C]. Output: ``N``. Caller does the Reduce::Mean` if they want a scalar.
Numerically stable (max-subtract before exp). Used for label
smoothing and knowledge distillation where targets are full
probability rows rather than class indices.
SoftmaxCrossEntropyWithLogits
Fused softmax + cross-entropy loss with integer (f32-encoded)
targets — the standard classification loss. Per-row output:
loss[n] = -log(softmax(logits[n])[labels[n]]).
Inputs: [logits, labels] with logits [N, C] and
labels [N] (f32-encoded class indices). Output: [N].
Caller does the Reduce::Mean if they want a scalar.
SoftmaxCrossEntropyBackward
Backward of the fused loss above. Emits
dlogits[n,c] = (softmax(logits[n])[c] - one_hot(labels)[n,c]) * d_loss[n].
Inputs: [logits, labels, d_loss]. Output: [N, C] (same shape
as logits). Recomputes the softmax inline rather than threading
it through from the forward node.
AttentionBackward
Backward of Op::Attention. Recomputes scaled QK^T, applies
the same mask_kind as the forward op, softmaxes scores, then
emits one of dQ, dK, or dV selected by AttentionBwdWrt.
Autodiff emits three nodes (one per wrt) so each output shape
stays a normal single-output MIR node.
Inputs: [q, k, v, dy] plus optional mask when mask_kind is
MaskKind::Custom or MaskKind::Bias (same convention as
forward). Output shape matches q, k, or v respectively.
AttentionBackwardAll
Fused backward of Op::Attention computing dQ, dK and dV from a
single score+softmax recompute (the three Op::AttentionBackward
siblings share q/k/v/dy, so recomputing scores 3× is wasteful). Created
by a backend-gated fusion pass (OpKind::AttentionBackwardAll). The
packed output is [3·B, S, H, D] (dQ‖dK‖dV stacked on the batch axis);
three axis-0 Narrows recover the individual gradients. Self-attention
only (q_seq == k_seq). Inputs match Op::AttentionBackward:
[q, k, v, dy] plus an optional mask.
FusedMatMulBiasAct
Fused matmul + bias + activation. Created from MatMul → Add → Activation.
Fields
activation: Option<Activation>FusedMatMulResidual
Fused matmul + full-tensor residual add: matmul(lhs, rhs) + residual,
single output (same shape as the matmul result). Created from
Add(MatMul(a, b), residual) when the MatMul feeds only the Add — the
transformer’s add(skip, o_proj) / add(h, down_proj) residuals. Unlike
Op::FusedMatMulBiasAct the third input is a full [m, n] activation
tensor (the residual/skip), not a rank-1 broadcast bias.
Emitted only for backends that implement the fused epilogue (today:
Metal, which folds the add into the sgemm store to save one dispatch on
the launch-latency-bound decode). Every other backend never runs the
fusion pass (gated by supports_op) and decomposes it in unfuse back
to MatMul + Add, so emitting it is always safe.
Inputs: [lhs, rhs, residual]. Output: [m, n].
FusedConvBiasAct
Fused convolution + bias + activation. Created from
Conv → Reshape(bias→[1,C,1,1]) → Expand → Add → [Activation].
Inputs: [input, weight, bias] — bias is rank-1 [C_out], mirroring
Op::QConv2d’s bundled-bias convention (not a separate broadcast Add).
Fields mirror Op::Conv plus the epilogue activation.
Only backends that can apply the whole epilogue in one dispatch claim
this op — today CUDA. Every other backend decomposes it back to
primitives in unfuse, so the fusion is always safe to emit. On CUDA it
lowers to cuDNN’s fused cudnnConvolutionBiasActivationForward. The
fusion pass only emits it for cuDNN-friendly shapes with an
identity/ReLU epilogue (the combination that benchmarks faster than
unfused); the conv_bias_act_epilogue kernel exists only as a
correctness fallback when libcudnn is absent at runtime.
has_residual: when true a 4th input residual (full [N,C,H,W] tensor)
is added before the activation — act(conv + bias + residual) — mapping
to cuDNN’s z operand. Used for ResNet conv→bias→add(residual)→relu.
Fields
activation: Option<Activation>FusedResidualLN
Fused residual + optional bias + layer norm. Created from Add(x, residual) → [Add(bias)] → LayerNorm.
FusedResidualRmsNorm
Fused residual + optional bias + RMS norm. Created from Add(x, residual) → [Add(bias)] → RmsNorm.
FusedSwiGLU
Fused SwiGLU: split input into up/gate halves, silu(gate) * up. Created from Split → Silu → Mul when fed by a concatenated matmul.
cast_to: optional output dtype — when Some(dt) the kernel casts
its result from the input dtype to dt in-register, saving a
separate cast pass. Reserved for future fp8/fp4 quantization paths;
for f32→f16 mixed precision the AutoMixedPrecision pass already
inserts a Cast node so this stays None in current pipelines.
Fields
FusedTransformerLayer
Fused full transformer layer: attention block + residual+LN + FFN + residual+LN. All intermediates resident in registers/threadgroup memory; one kernel per layer instead of ~30 (the CPU’s batch=1 win, lifted to IR so any backend can implement it as a monolithic kernel).
Inputs: hidden, qkv_w, qkv_b, out_w, out_b, ln1_g, ln1_b, fc1_w, fc1_b, fc2_w, fc2_b, ln2_g, ln2_b, mask Output: same shape as hidden.
Backend status: same as FusedAttentionBlock. CPU implements the L1-cache-resident merge at the thunk level. Metal deferred — requires a single MSL kernel for the whole layer to actually beat the unfused path. Multi-day work; revisit when there’s a model whose Metal inference is bottlenecked here rather than on the wait latency floor.
Fields
activation: ActivationFusedAttentionBlock
Fused attention block: QKV projection → split → [RoPE] → SDPA → output projection. Created by FuseAttentionBlock pass when batch*seq is small. All intermediates stay in L1 cache — no arena writes between ops.
Inputs (in order): hidden, qkv_w, out_w, mask, [qkv_b, out_b] if has_bias, [rope_cos, rope_sin] if has_rope
Backend status (Phase C finalize):
CPU — implemented at the thunk level: the CPU schedule
recognizes the multi-thunk pattern and merges into
a single FusedAttnBlock that keeps Q/K/V in stack
buffers across stages (the L1-cache win).
Metal — deferred. A dispatch-wrapper version (chaining
existing kernels) buys nothing the unfused Metal path
doesn’t already get, since per-run cost is dominated
by wait_until_completed (~150 µs), not encode. The
real win is a single MSL kernel keeping Q/K/V in
threadgroup memory across stages — multi-day work.
Until then, Metal runs the unfused chain (one matmul,
three narrows, two ropes, attention, one matmul) — all
covered in op_coverage and parity_harness.
If
Conditional: pick between two subgraphs based on a boolean predicate.
Inputs: [predicate, …captures (used inside both branches)].
then_branch and else_branch are sub-graphs that share the
captured inputs and must produce identically-shaped outputs.
Used for: shape-dependent execution, batched inference of
dynamic-length sequences with padding masks.
While
Loop: iterate body while cond evaluates true.
Inputs: […initial loop-carried values].
cond’s single output is a Bool scalar.
body’s outputs become the next iteration’s loop-carried inputs.
Outputs of While are the values after the final iteration.
Used for: KV-cache-driven autoregressive generation, beam search.
Scan
Bounded-length loop with a fixed-shape carry, optional per-step
inputs, and optional stacked output. Mirrors JAX’s lax.scan.
Body signature: (carry, x_t_0, ..., x_t_{num_xs-1}) → carry_next
— 1 + num_xs Op::Inputs in NodeId construction order (first
declared is the carry; the remaining num_xs are per-step
slices). Single output (the next carry).
Outer Op::Scan inputs (in order):
[init_carry, xs_0, xs_1, ..., xs_{num_xs-1}]
Each xs_i has shape [length, *per_step_shape_i]; the body
sees xs_i[t] (a per_step_shape_i slice) on iteration t.
Outer Op::Scan output:
save_trajectory == false— final carry, shape*carry.save_trajectory == true— stacked trajectory of carries, shape[length, *carry]. Rowtis the carry after stept+1, so rowlength-1matches the no-trajectory case.
Mirrors JAX’s lax.scan. Common uses include time-stepping
integrators with time-varying drives, Mamba-style SSM scans
reading per-step inputs, and RNN-style sequence processing.
Fields
num_bcast: u32Number of “broadcast” inputs — values that are constant
across iterations. Outer scan inputs in order:
[init, bcast_0..bcast_{B-1}, xs_0..xs_{X-1}]
Body Op::Inputs in NodeId order:
[carry, bcast_0..bcast_{B-1}, x_t_0..x_t_{X-1}]
CPU executor fills bcast slots ONCE before the iteration
loop (xs slots are filled per-step). The reverse-mode AD
pre-pass materialises each bcast into an xs of shape
[length, *bcast] via broadcast Mul so the rest of the
VJP / executor pipeline can stay unchanged. 0 (default)
keeps the original carry+xs scan shape.
num_checkpoints: u32Number of trajectory checkpoints when save_trajectory == true. 0 means “save all length rows” (default). A
positive value K means save only K evenly-spaced rows
at indices floor(t * length / K) for t in 0..K. Used
by recursive checkpointed AD: store O(√T) carries during
forward, recompute the rest in the backward pass.
When 0 (or K == length), the saved trajectory has
shape [length, *carry] — same as the original behavior.
When 0 < K < length, the saved trajectory has shape
[K, *carry].
ScanBackward
Reverse-mode AD companion to Op::Scan — extracts the carry
gradient dinit. Walks t = length-1 .. 0, applying body_vjp
to thread dcarry back through the time loop.
Inputs (in order):
[init, trajectory, upstream, xs_0, ..., xs_{num_xs-1}]
Output: dinit, shape = carry shape.
body_vjp is the result of
autodiff::grad(body, [carry_id, xs_0_id, ..., xs_{num_xs-1}_id])
— a graph with 1 + num_xs + 1 Inputs (carry + x_t_i for each
xs + "d_output") and 1 + num_xs outputs
(dcarry + dx_t_i for each xs). This op reads outputs[0] =
dcarry; the sibling Self::ScanBackwardXs reads the
outputs[1 + xs_idx] slot for each xs gradient.
Fields
num_checkpoints: u32When 0 or equal to length, the trajectory input has
shape [length, *carry] — every step’s carry is cached
(CheckpointStrategy::All). When 0 < K < length, the
trajectory input has shape [K, *carry] and the executor
recomputes intermediate carries via forward_body between
checkpoints. forward_body must be Some whenever this
is < length.
ScanBackwardXs
Companion to Self::ScanBackward that extracts one stacked
per-step dxs_i (shape [length, *per_step_xs_i]). Same inputs
and same body_vjp graph as ScanBackward — xs_idx selects
which body_vjp output to stack into the result.
Note: each ScanBackwardXs runs its own backward loop. A future
optimization can fuse them into a single multi-output backward
kernel; for now it’s 1 + num_xs independent sweeps.
Fields
GaussianSplatRender
CPU reference 3D Gaussian splat forward render.
Seven flat F32 inputs (scene buffers + camera/render meta):
0. positions [N*3]
- scales
[N*3](log-space) - rotations
[N*4](xyzw) - opacities
[N](logit) - colors
[N*3](linear RGB) - sh_coeffs
[N * sh_coeff_count * 3] - meta
[23]— camera position/target/up/fov/near/far, background RGB, then width/height/tile_size/radius_scale/alpha_cutoff/max_splat_steps/ transmittance_threshold/max_list_entries as f32 bit-patterns.
Output: [height * width * 4] linear RGBA (display gamma baked in).
Build via crate::Graph::gaussian_splat_render.
Differentiable backward is not implemented in v1; autodiff treats this
op as non-differentiable (same as Op::Sample).
Fields
GaussianSplatRenderBackward
Backward pass for Self::GaussianSplatRender.
Eight inputs: the same seven as forward plus d_loss_rgba [W*H*4]
(only RGB channels are used). Re-runs the training forward internally.
Output: packed gradients
[positions(3N) | scales(3N) | rotations(4N) | opacities(N) | colors(3N) | sh(N*sh*3)].
Unpack via crate::ops::splat::unpack_gaussian_splat_packed_grads.
Fields
GaussianSplatPrepare
Strict IR stage 1: project, bin, sort, build per-pixel rays.
Seven inputs (same scene + meta as Self::GaussianSplatRender). Output: packed
prepare buffer (see rlx_splat::prep_layout::prep_packed_len).
Fields
GaussianSplatRasterize
Strict IR stage 2: tile raster from Self::GaussianSplatPrepare output.
Inputs: prep packed buffer, meta [23]. Output: [width * height * 4] RGBA.
Fields
Custom
User-registered custom op. name keys into the
crate::op_registry for shape inference, autodiff, and
per-backend execution. attrs is an opaque blob passed
through to those callbacks (FFT direction, SparseLU
reordering strategy, etc.). num_inputs is captured at
construction time so Op::num_inputs stays infallible
without a registry lookup. Build via crate::Graph::custom_op.
Fft
1D Fast Fourier Transform along the last axis.
Layouts
F32/F64: 2N real-block — last axis is[re₀…re_{N-1}, im₀…im_{N-1}].C64: interleaved[re, im]pairs per complex element along the last axis.
ND transforms — use Graph::fftn / Graph::ifftn, which compose
fft_axis (transpose → Fft → transpose). Multi-axis fftn requires
DType::C64; the 2N-block layout describes a single complex axis.
Default (FftNorm::Backward) is unnormalized on both directions:
fft(x)[k] = Σ x[n]·exp(-2πi·nk/N)
ifft(y)[n] = Σ y[k]·exp(+2πi·nk/N)
so ifft(fft(x)) = N·x. Use FftNorm::Forward for gpu-fft-style
1/N scaling on inverse, or FftNorm::Ortho for unitary scaling.
AD: VJP(fft) = ifft, VJP(ifft) = fft when norm=Backward;
other norms apply the chain rule via output scaling.
FftQ
Fixed-point FFT over I32 data in the same 2N-block layout Op::Fft
uses for F32 (n real parts then n imaginary parts per row).
Separate from Op::Fft rather than a dtype of it because it needs a
parameter the floating-point transform does not: how to keep the
datapath in range. That choice costs precision, so it is explicit —
see crate::fft::FftQScale.
Radix-2 only; the length must be a power of two.
FftButterflyStage
Ternary pruned radix-2 butterfly stage on interleaved complex state.
Inputs:
0 — state [batch, n_fft, 2] (re/im on axis 2)
1 — gate [half] — 0 = identity, 1 = run butterfly (half = n_fft/2)
2 — rev [half] — 0 = forward, 1 = swap outputs when gate=1
3 — tw_re [half]
4 — tw_im [half]
Output: [batch, n_fft, 2] same layout. Slots with gate=0 copy inputs
without twiddle math.
LogMel
Log-mel spectrogram from RLX FFT block-layout spectrum.
Inputs:
0 — spectrum [..., 2*n_fft] (re plane then im plane, same as Op::Fft output)
1 — mel filterbank [n_mels, n_bins] with n_bins = n_fft/2 + 1
Output: [..., n_mels] with Whisper dynamic-range compression
(log10, clamp to max−8 dB, (x+4)/4).
LogMelBackward
VJP of Op::LogMel w.r.t. spectrum (input 0).
Inputs: spectrum block, mel filters, upstream dy.
Output: d_spectrum (same shape as input 0).
WelchPeaks
Top-K Welch peaks from block-layout segment spectra.
Input 0: spectrum [batch * n_segments, 2*n_fft] (re ∥ im planes).
Output: [batch, k*2] interleaved (bin, power) per spike.
CustomFn
User-defined sub-graph with optional override AD rules.
Mirrors JAX’s custom_vjp / custom_jvp decorators: the
caller wraps a forward computation and supplies its own
reverse- and/or forward-mode AD bodies. Useful when:
- The forward is iterative (Newton, fixed-point) and differentiating through the loop is wasteful — the vjp_body computes the implicit-function gradient at the converged point in one shot.
- The math has a closed-form gradient that’s much cheaper than autodiff.
- The forward op is non-differentiable by tracing (sampling, argmax) and the user wants a smooth surrogate.
fwd_body: num_inputs Op::Inputs in NodeId construction
order, one Op::Output (the primal y). Forward execution
inlines this body once.
vjp_body (optional): Op::Inputs are num_inputs primal
inputs in NodeId order, plus two special-named Inputs —
"primal_output" (the y from forward) and "d_output" (the
upstream gradient). Outputs: num_inputs tensors in
set_outputs order, matching the gradients of each primal
input. When None, reverse-mode AD recurses into fwd_body
— same as if the op were inlined.
jvp_body (optional): Op::Inputs are num_inputs primal
inputs in NodeId order, num_inputs special-named Inputs
"tangent_0"..="tangent_{num_inputs-1}" carrying each input’s
tangent, and an optional special-named "primal_output" Input
(the y from forward, useful when the JVP must be evaluated at
a converged / nonlinear point — e.g. IFT-style forward-mode AD
of an iterative solver). Output: 1 tensor (the tangent of y).
When None, forward-mode AD recurses into fwd_body.
num_inputs is captured so Op::num_inputs stays
infallible. Build via crate::Graph::custom_fn.
Fields
BiMap
BiMap (bilinear mapping) SPDNet layer: Y = W · X · Wᵀ.
Inputs: [W [m,n], X [n,n]] (X symmetric SPD). Output: Y [m,m]
(SPD when W has full row rank). Reduces the manifold dimension.
The Stiefel / semi-orthogonality constraint on W is the
optimizer’s responsibility, not this op’s. VJP decomposes to
matmuls: dW = (Ḡ+Ḡᵀ)·W·X, dX = Wᵀ·Ḡ·W.
ReEig
ReEig (eigenvalue rectification) SPDNet nonlinearity:
Y = U · max(ε, Σ) · Uᵀ where X = U Σ Uᵀ. Input [X [n,n]]
symmetric → Y [n,n] SPD. The SPD analogue of ReLU: floors the
spectrum at eps to keep the output well-conditioned.
LogEig
LogEig SPDNet layer: Y = U · log(Σ) · Uᵀ = logm(X). Maps the
SPD manifold to the tangent space at the identity (symmetric
matrices) so a plain Euclidean head can classify. Input
[X [n,n]] → Y [n,n]. eps floors the spectrum before log.
SpdBatchNorm
SPD batch-norm affine transport (Brooks et al. 2019):
Y_i = G^{1/2} · (M^{-1/2} X_i M^{-1/2}) · G^{1/2}.
Inputs: [X [B,n,n], M [n,n], G [n,n]] — a batch of SPD matrices,
the (detached) Fréchet mean M, and the learnable SPD bias G.
Output: Y [B,n,n]. Centers the batch to the identity then biases
toward G. M receives no gradient (detached statistic, like
Euclidean batch-norm’s running mean); grads flow to X and G.
SpdKarcherMean
Fréchet / Karcher mean of a batch of SPD matrices under the AIRM
metric (iterative tangent averaging). Input [X [B,n,n]] →
M [n,n]. Non-differentiable (a detached batch statistic — its
VJP contributes no gradient), mirroring how batch-norm’s running
statistics live outside the autodiff graph.
SpdKarcherMeanWeighted
Weighted Fréchet / Karcher mean of a batch of SPD matrices under the
AIRM metric — the barycentre argmin_M Σᵢ wᵢ · δ²(M, Cᵢ). Inputs
[X [B,n,n], w [B]] → M [n,n]. The weighted counterpart of
Op::SpdKarcherMean: the true AIRM barycentre a barycentric-OT
projection / soft-clustering / weighted-MDM needs (not the
log-Euclidean shortcut expm(Σ w̄ᵢ logm Cᵢ)). Non-differentiable — a
detached statistic, exactly like Op::SpdKarcherMean.
SpdLogMap
AIRM Riemannian logarithm at an arbitrary base point:
Log_P(X) = P^{1/2} · logm(P^{-1/2} X P^{-1/2}) · P^{1/2}. Inputs
[base [n,n], x [n,n]] → tangent vector [n,n]. Generalises
Op::LogEig (which is Log_I) to any base point — the Riemannian
OT / domain-adaptation primitive (Yair et al. 2019).
SpdExpMap
AIRM Riemannian exponential at an arbitrary base point:
Exp_P(V) = P^{1/2} · expm(P^{-1/2} V P^{-1/2}) · P^{1/2}. Inputs
[base [n,n], v [n,n]] → SPD point [n,n]. Inverse of
Op::SpdLogMap.
SpdParallelTransport
AIRM parallel transport of a tangent vector between base points:
Γ_{P→Q}(V) = E · V · Eᵀ, E = (Q P^{-1})^{1/2}. Inputs
[from [n,n], to [n,n], v [n,n]] → transported tangent vector [n,n].
The AIRM isometry that moves tangent vectors between base points for
Riemannian OT / domain adaptation (Yair et al. 2019).
SpdMatrixFnBatch
Batched SPD spectral matrix function — applies kind (logm / expm /
sqrtm / invsqrtm) to each matrix of a batch. Input [X [B,n,n]] →
[B,n,n]. The graph-level, any-backend counterpart of
rlx_cpu::spd::{logm,expm,sqrtm,invsqrtm}_batch.
ReEigBackward
Backward of Op::ReEig: (X [n,n], dY [n,n]) → dX [n,n] via
the Daleckii–Krein (Loewner) adjoint with f(λ)=max(λ,ε).
LogEigBackward
Backward of Op::LogEig: (X, dY) → dX via the Loewner adjoint
with f=log, f′=1/λ.
SpdBatchNormBackwardX
Backward of Op::SpdBatchNorm w.r.t. the inputs X (mean
detached): (M [n,n], G [n,n], dY [B,n,n]) → dX [B,n,n].
SpdBatchNormBackwardG
Backward of Op::SpdBatchNorm w.r.t. the learnable bias G
(mean detached): (X [B,n,n], M [n,n], G [n,n], dY [B,n,n]) → dG [n,n].
SpdLogMapBackward
Backward of Op::SpdLogMap. Inputs [base [n,n], x [n,n], dY [n,n]] →
packed [d_base ∥ d_x] of shape [2, n, n] (the builder narrows the two
gradients out). Full AIRM adjoint — the base point receives a gradient.
SpdExpMapBackward
Backward of Op::SpdExpMap. Inputs [base [n,n], v [n,n], dY [n,n]] →
packed [d_base ∥ d_v] of shape [2, n, n].
SpdParallelTransportBackward
Backward of Op::SpdParallelTransport.
Inputs [from [n,n], to [n,n], v [n,n], dY [n,n]] → packed
[d_from ∥ d_to ∥ d_v] of shape [3, n, n].
SpdMatrixFnBatchBackward
Backward of Op::SpdMatrixFnBatch: per-slice Daleckii–Krein adjoint of
kind. Inputs [X [B,n,n], dY [B,n,n]] → dX [B,n,n].
Eigh
Symmetric eigendecomposition of an [n,n] matrix. Output is packed
[λ (n) ∥ U (n²)] — ascending eigenvalues then the row-major eigenvector
matrix (column j = eigenvector j, so A = U diag(λ) Uᵀ). The builder
Graph::eigh narrows the two views out. Differentiable (see
Op::EighBackward). The seam a native batched eigensolver (cuSOLVER
syevjBatched &c.) plugs into.
EighBackward
Backward of Op::Eigh. Inputs [fwd [n²+n], bar [n²+n]] — the packed
forward [λ ∥ U] and the upstream cotangent [λ̄ ∥ Ū] — → Ā [n,n]
(the standard symmetric-eigendecomposition adjoint).
EighBatch
Batched Op::Eigh: [B,n,n] → [B, n²+n] packed per slice.
EighBatchBackward
Backward of Op::EighBatch. Inputs [fwd [B,n²+n], bar [B,n²+n]] →
Ā [B,n,n].
AdaLayerNorm
Fused adaptive-norm modulation — the signature DiT / adaLN-Zero op.
out = norm(x) * (1 + scale) + shiftnorm normalizes over the last axis with no learnable affine
(AdaNormKind); the affine is supplied per-conditioning by scale
and shift, which broadcast over every axis except the last (the DiT
[B, 1, D] modulation over [B, S, D]). Folding the broadcast into the
op avoids the Op::Expand materialization the decomposed form needs
([B,1,D]→[B,S,D] can’t ride an ElementwiseRegion’s trailing-modulo
broadcast) and keeps the pre-norm activation off the heap — mirroring
Op::FusedResidualLN.
Inputs: [x, scale, shift]
x:[..., D]scale: broadcast-compatible withx, last dimD(e.g.[B, 1, D])shift: same shape asscale
Output shape == x. Backends without a native kernel decompose via
unfuse into norm → mul(1+scale) → add(shift).
GatedResidual
Gated residual out = x + gate * y (the DiT post-sublayer combine).
gate broadcasts over every axis except the last (the [B, 1, D]
modulation gate over [B, S, D]), so — like Op::AdaLayerNorm — the
built-in broadcast avoids an Op::Expand of the gate. Used for both the
attention (gate_msa) and MLP (gate_mlp) residual joins.
Inputs: [x, y, gate]
x: residual stream[..., D]y: sublayer output, same shape asxgate: broadcast-compatible withx, last dimD(e.g.[B, 1, D])
Output shape == x. Backends without a native kernel decompose via
unfuse into mul(gate, y) → add(x).
AdaLayerNormBackward
Packed reverse of Op::AdaLayerNorm.
Inputs: [x, scale, shift, dy] (same shapes as the forward).
Output: 1-D [nx + 2·ns] packing dx (numel nx), dscale (ns),
dshift (ns) in that order — avoids materializing Expand of
[B,1,D] modulation in the backward graph. Unpack with
crate::ops::backward::unpack_ada_layer_norm_backward.
GatedResidualBackward
Packed reverse of Op::GatedResidual.
Inputs: [x, y, gate, dy]. Output: 1-D [nx + ny + ng] packing
dx, dy, dgate (row-major). Unpack with
crate::ops::backward::unpack_gated_residual_backward.
Implementations§
Source§impl Op
impl Op
Sourcepub fn is_elementwise(&self) -> bool
pub fn is_elementwise(&self) -> bool
True if this op is element-wise (same shape in, same shape out). Element-wise ops are prime fusion candidates.
Sourcepub fn is_transform_eligible(&self) -> bool
pub fn is_transform_eligible(&self) -> bool
True if this op may appear in a Op::TransformRegion chain.
Sourcepub fn is_blas(&self) -> bool
pub fn is_blas(&self) -> bool
True if this op is a BLAS/compute-intensive op that forms a fusion boundary.
Sourcepub fn is_fusion_boundary(&self) -> bool
pub fn is_fusion_boundary(&self) -> bool
True if element-wise fusion must not span across this op.
Sourcepub fn is_reduction(&self) -> bool
pub fn is_reduction(&self) -> bool
True if this op is a reduction (drives loop iteration in fused kernels).
Sourcepub fn is_const_foldable(&self) -> bool
pub fn is_const_foldable(&self) -> bool
True if the constant folder can evaluate this op on the host.
Sourcepub fn is_fused(&self) -> bool
pub fn is_fused(&self) -> bool
True if this is a composite op with a decomposition into primitives.
Sourcepub fn has_nested_body(&self) -> bool
pub fn has_nested_body(&self) -> bool
True if this op carries nested Graph bodies.
Equivalent to !self.subgraphs().is_empty(), without the walk.
Sourcepub fn is_nondeterministic(&self) -> bool
pub fn is_nondeterministic(&self) -> bool
True if the output is not a function of the inputs alone (RNG / sampling), so it must never be folded, CSE’d or hoisted.
Source§impl Op
impl Op
Sourcepub fn subgraphs(&self) -> Vec<&Graph>
pub fn subgraphs(&self) -> Vec<&Graph>
Every nested Graph this op carries, in field order.
Empty for all but the control-flow / custom-body ops
(Op::If, Op::While, Op::Scan, Op::ScanBackward,
Op::ScanBackwardXs, Op::CustomFn). The empty Vec does not
allocate, so the common path is free.
Sourcepub fn subgraphs_mut(&mut self) -> Vec<&mut Graph>
pub fn subgraphs_mut(&mut self) -> Vec<&mut Graph>
Mutable counterpart of Op::subgraphs, same order.
Sourcepub fn kind(&self) -> OpKind
pub fn kind(&self) -> OpKind
PLAN L4: discriminant for backend-supported-set checks.
Stable, parameter-free identity per variant — Op::Activation(_)
and Op::Activation(Relu) share the same OpKind::Activation.
Sourcepub fn arity(&self) -> Arity
pub fn arity(&self) -> Arity
Operand count this op accepts.
The single source of truth for arity. Prefer this over
num_inputs, which collapses a range to its
minimum and so cannot distinguish “exactly 2” from “2 or 3”.
Sourcepub fn num_inputs(&self) -> usize
pub fn num_inputs(&self) -> usize
Minimum operands this op requires.
Kept for callers that just need a count. For ops with an optional
operand this is the lower bound, so it cannot be used to reject an
operand list — use arity and Arity::accepts.
Trait Implementations§
Source§impl<'de> Deserialize<'de> for Op
impl<'de> Deserialize<'de> for Op
Source§fn deserialize<__D>(
__deserializer: __D,
) -> Result<Op, <__D as Deserializer<'de>>::Error>where
__D: Deserializer<'de>,
fn deserialize<__D>(
__deserializer: __D,
) -> Result<Op, <__D as Deserializer<'de>>::Error>where
__D: Deserializer<'de>,
Source§impl Serialize for Op
impl Serialize for Op
Source§fn serialize<__S>(
&self,
__serializer: __S,
) -> Result<<__S as Serializer>::Ok, <__S as Serializer>::Error>where
__S: Serializer,
fn serialize<__S>(
&self,
__serializer: __S,
) -> Result<<__S as Serializer>::Ok, <__S as Serializer>::Error>where
__S: Serializer,
impl StructuralPartialEq for Op
Auto Trait Implementations§
impl Freeze for Op
impl RefUnwindSafe for Op
impl Send for Op
impl Sync for Op
impl Unpin for Op
impl UnsafeUnpin for Op
impl UnwindSafe for Op
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Source§impl<T> CloneToUninit for Twhere
T: Clone,
impl<T> CloneToUninit for Twhere
T: Clone,
impl<T> DeserializeOwned for Twhere
T: for<'de> Deserialize<'de>,
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ
fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ
self into a Left variant of Either<Self, Self>
if into_left is true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
self into a Left variant of Either<Self, Self>
if into_left(&self) returns true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read more