Enum OpKind 

pub enum OpKind {
    Constant,
    Parameter,
    Add,
    Sub,
    Mul,
    Div,
    Neg,
    Exp,
    Log,
    Sum {
        axis: Option<i32>,
    },
    Mean {
        axis: Option<i32>,
    },
    Max {
        axis: Option<i32>,
    },
    MatMul,
    Reshape {
        new_shape: Shape,
    },
    Transpose {
        axes: Option<Vec<usize>>,
    },
    Softmax {
        axis: i32,
    },
    Silu,
    Gelu,
    LayerNorm {
        eps: f32,
    },
    RmsNorm {
        eps: f32,
    },
    Rope {
        rotary_dim: usize,
        pos_offset: usize,
        theta: f32,
    },
    Broadcast {
        target_shape: Shape,
    },
    ScaledMaskedSoftmax {
        scale: f32,
        causal: bool,
    },
    Attention {
        scale: f32,
        causal: bool,
    },
    LayerNormVjp {
        eps: f32,
    },
    RmsNormVjp {
        eps: f32,
    },
    SoftmaxVjp {
        axis: i32,
    },
    SiluVjp,
    GeluVjp,
    Sqrt,
    RoPE {
        base: f32,
        offset: usize,
        traditional: bool,
    },
    Embedding,
    Narrow {
        axis: i32,
        start: i64,
        length: i64,
    },
    Concatenate {
        axis: i32,
    },
}

The set of operations supported by the graph IR.

Variants

Constant

Constant tensor (data already materialized).

Parameter

Parameter (learnable weight, data provided externally).

Add

Sub

Mul

Div

Neg

Exp

Log

Sum

Fields

axis: Option<i32>

Mean

Fields

axis: Option<i32>

Max

Fields

axis: Option<i32>

MatMul

Reshape

Fields

new_shape: Shape

Transpose

Fields

axes: Option<Vec<usize>>

Softmax

Fields

axis: i32

Silu

Gelu

LayerNorm

Fields

eps: f32

RmsNorm

Fields

eps: f32

Rope

Rotary positional embeddings (RoPE). Applied in-place to interleaved pairs of the last dimension.

Fields

rotary_dim: usize

pos_offset: usize

theta: f32

Broadcast

Broadcast a tensor to a target shape (numpy-style rules).

Fields

target_shape: Shape

ScaledMaskedSoftmax

Fused scale + causal-mask + softmax along last axis. Input: scores [Tq, Tk], output: probs [Tq, Tk]

Fields

scale: f32

causal: bool

Attention

Full single-head attention composition. Inputs: [Q, K, V] where Q=[Tq,Dh], K=[Tk,Dh], V=[Tk,Dh] Output: Y=[Tq,Dh]

Fields

scale: f32

causal: bool

LayerNormVjp

LayerNorm backward: inputs = [grad_output, input], produces grad_input.

Fields

eps: f32

RmsNormVjp

RmsNorm backward: inputs = [grad_output, input], produces grad_input.

Fields

eps: f32

SoftmaxVjp

Softmax backward: inputs = [grad_output, softmax_output], produces grad_input.

Fields

axis: i32

SiluVjp

SiLU backward: inputs = [grad_output, original_input], produces grad_input.

GeluVjp

GELU backward: inputs = [grad_output, original_input], produces grad_input.

Sqrt

Element-wise square root.

RoPE

Fields

base: f32

offset: usize

traditional: bool

Embedding

Embedding lookup: gather rows from a weight matrix by index. Inputs: [weight [vocab, dim], indices [seq_len]] Output: [seq_len, dim]

Narrow

Extract a contiguous slice along an axis. Inputs: [input] Output: narrowed tensor

Fields

axis: i32

start: i64

length: i64

Concatenate

Concatenate tensors along an axis. Inputs: [tensor_0, tensor_1, …] Output: concatenated tensor

Fields

axis: i32

Trait Implementations

impl Clone for OpKind

fn clone(&self) -> OpKind

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for OpKind

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

Formats the value using the given formatter.