Enum FastOp 

#[non_exhaustive]
pub enum FastOp {
    RmsNorm {
        eps: u32,
    },
    AddRmsNorm {
        eps: u32,
    },
    RmsNormVjp {
        eps: u32,
    },
    LayerNorm {
        eps: u32,
    },
    LayerNormVjp {
        eps: u32,
    },
    Rope {
        dims: usize,
        traditional: bool,
        base: u32,
        scale: u32,
        forward: bool,
    },
    NormRope {
        eps: Option<u32>,
        dims: usize,
        traditional: bool,
        base: u32,
        scale: u32,
        offset: usize,
        cache: bool,
    },
    DepthwiseConv1d {
        stride: usize,
        padding: usize,
        dilation: usize,
    },
    Softmax,
    LogSumExp,
    CrossEntropy,
    CrossEntropyVjp,
    Attention {
        scale: u32,
        causal: bool,
        mask: bool,
        sinks: bool,
    },
    GatedDelta {
        steps: bool,
    },
    GatherMm,
    GatherQmm {
        group_size: usize,
        bits: u8,
        transpose: bool,
    },
}

A fused operation (see the module documentation).

Variants (Non-exhaustive)

Non-exhaustive enums could have additional variants added in future. Therefore, when matching against variants of non-exhaustive enums, an extra wildcard arm must be added to account for any future variants.

RmsNorm

(x, w): x / sqrt(mean(x²) + eps) · w over the last axis; w is 1-D or a scalar 1.

Fields

eps: u32

f32::to_bits of epsilon.

AddRmsNorm

(x, r, w): h = x + r and FastOp::RmsNorm of h (see add_rms_norm).

Fields

eps: u32

f32::to_bits of epsilon.

RmsNormVjp

(x, w, g): the gradients of FastOp::RmsNorm for x and w.

Fields

eps: u32

f32::to_bits of epsilon.

LayerNorm

(x, w, b): (x - mean) / sqrt(var + eps) · w + b over the last axis; w and b are 1-D or the scalars 1 and 0.

Fields

eps: u32

f32::to_bits of epsilon.

LayerNormVjp

(x, w, b, g): the gradients of FastOp::LayerNorm for x, w and b.

Fields

eps: u32

f32::to_bits of epsilon.

Rope

(x, offset[, freqs]): rotary position embeddings of the first dims features.

Fields

dims: usize

Rotated features (even).

traditional: bool

Rotate interleaved pairs instead of the two halves.

base: u32

f32::to_bits of the frequency base (unused with freqs).

scale: u32

f32::to_bits of the position scale.

forward: bool

Rotate forwards (false: the inverse rotation, the gradient).

NormRope

([cache,] x, w[, freqs]): norm_rope: each [b, h, l] row of x RMS-normalised (with eps) and rotated as FastOp::Rope at position offset + l; written into cache at the same positions, or returned as a new [B, H, L, D] array.

Fields

eps: Option<u32>

f32::to_bits of epsilon, or no normalisation.

dims: usize

Rotated features (even).

traditional: bool

Rotate interleaved pairs instead of the two halves.

base: u32

f32::to_bits of the frequency base (unused with freqs).

scale: u32

f32::to_bits of the position scale.

offset: usize

The position of the first row.

cache: bool

Whether the first input is a cache written at the rows’ positions.

DepthwiseConv1d

(x, w): crate::conv1d with groups equal to the channels (depthwise).

Fields

stride: usize

Positions between outputs.

padding: usize

Zeros before and after the input.

dilation: usize

Positions between taps.

Softmax

(x): softmax over the last axis, accumulated in float32.

LogSumExp

(x): log(sum(exp(x))) over the last axis, which becomes size 1.

CrossEntropy

(logits, targets): logsumexp(logits) - logits[target] per row, in float32.

CrossEntropyVjp

(logits, targets, loss, g): the gradient of FastOp::CrossEntropy for the logits.

Attention

(q, k, v[, mask][, sinks]): scaled dot-product attention (see scaled_dot_product_attention).

Fields

scale: u32

f32::to_bits of the score scale.

causal: bool

Causal masking (keys after the query’s position, aligned at the end, are hidden).

mask: bool

Whether an array mask (bool or additive, [B, Hq, Lq, Lk]) is the fourth input.

sinks: bool

Whether per-head attention sinks ([Hq]) are the last input.

GatedDelta

(q, k, v, g, beta, state): the gated delta rule (see gated_delta), with the outputs and the final state (or every step’s, gated_delta_steps).

Fields

steps: bool

The state after every step ([B, T, Hv, Dv, Dk]) instead of the final one.

GatherMm

(a, b, lhs_indices, rhs_indices[, offsets]): matrix products of gathered batches (see crate::gather_mm), as FastOp::GatherQmm in full precision.

GatherQmm

(x, w, scales, biases, lhs_indices, rhs_indices[, offsets]): quantized matrix products of gathered batches (see crate::gather_qmm); the indices are uint32 and broadcast to the output’s batch shape. offsets (int32, one per batch of w) marks where each weight batch’s rows start when the rows are sorted by it.

Fields

group_size: usize

Values per scale and bias.

bits: u8

Bits per value.

transpose: bool

x · wᵀ (else x · w).

Trait Implementations

impl Clone for FastOp

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for FastOp

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Eq for FastOp

impl Hash for FastOp

fn hash<__H: Hasher>(&self, state: &mut __H)

Feeds this value into the given Hasher.

fn hash_slice<H>(data: &[Self], state: &mut H)

where H: Hasher, Self: Sized,

Feeds a slice of this type into the given Hasher.

impl PartialEq for FastOp

fn eq(&self, other: &Self) -> bool

Equality operator ==.

fn ne(&self, other: &Rhs) -> bool

Inequality operator !=.

impl StructuralPartialEq for FastOp