Struct RotaryEmbedding 

pub struct RotaryEmbedding {
    pub head_dim: usize,
    pub max_positions: usize,
    pub theta: f32,
    /* private fields */
}

Rotary Positional Embedding (RoPE).

Precomputes cos and sin tables for all positions up to max_positions. The rotation applies to pairs of dimensions (x_{2i}, x_{2i+1}) as:

x_out[2i]   = x[2i]*cos(θ_i*pos) − x[2i+1]*sin(θ_i*pos)
x_out[2i+1] = x[2i]*sin(θ_i*pos) + x[2i+1]*cos(θ_i*pos)

where θ_i = theta ^ (-2i / head_dim).

This embeds position information directly into the attention dot product without requiring separate positional embeddings.

Fields

head_dim: usize

Head dimension (must be even).

max_positions: usize

Maximum sequence length for which tables are precomputed.

theta: f32

RoPE base frequency (typically 10 000 for LLaMA-2, 500 000 for LLaMA-3).

Implementations

impl RotaryEmbedding

pub fn new(head_dim: usize, max_positions: usize, theta: f32) -> LmResult<Self>

Build RoPE tables for the given configuration.

pub fn apply( &self, x: &mut [f32], n_heads: usize, n_tokens: usize, offset: usize, ) -> LmResult<()>

Apply RoPE in-place to a QKV projection.

x has shape [n_tokens × n_heads × head_dim]. offset is the absolute position of the first token (for KV-cache decode).

pub fn cos_at(&self, pos: usize, i: usize) -> f32

Cosine value for (position, half_dim_index).

pub fn sin_at(&self, pos: usize, i: usize) -> f32

Sine value for (position, half_dim_index).

Trait Implementations

impl Clone for RotaryEmbedding

fn clone(&self) -> RotaryEmbedding

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for RotaryEmbedding

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

Formats the value using the given formatter.