Struct RoPE 

pub struct RoPE {
    pub cos_cache: DenseTensor,
    pub sin_cache: DenseTensor,
    pub dim: usize,
    pub max_seq_len: usize,
    pub base: f64,
}

Rotary Positional Embedding (RoPE)

RoPE encodes positional information by rotating the query and key vectors based on their positions. This is used in LLaMA, Mistral, and other modern LLMs.

Formula: RoPE(x, pos) = x * cos(pos * theta) + rotate_half(x) * sin(pos * theta)

where theta = base^(-2i/dim) for i = 0, 1, …, dim/2-1

Fields

cos_cache: DenseTensor

Pre-computed cosine cache [max_seq_len, dim/2]

sin_cache: DenseTensor

Pre-computed sine cache [max_seq_len, dim/2]

dim: usize

Rotation dimension (typically head_dim)

max_seq_len: usize

Maximum sequence length

base: f64

Base frequency (LLaMA uses 10000)

Implementations

impl RoPE

pub fn new(dim: usize, max_seq_len: usize, base: f64) -> Self

Create a new RoPE module

Arguments

• dim - Rotation dimension (typically head_dim)
• max_seq_len - Maximum sequence length to support
• base - Base frequency for theta calculation (default: 10000)

pub fn default(dim: usize, max_seq_len: usize) -> Self

Create RoPE with default base (10000)

pub fn forward( &self, x: &DenseTensor, positions: Option<&[usize]>, ) -> DenseTensor

Apply RoPE to input tensor

Arguments

• x - Input tensor [batch_size, seq_len, dim]
• positions - Optional position indices [seq_len] or [batch_size * seq_len]

Returns

Rotated tensor [batch_size, seq_len, dim]

pub fn forward_qk( &self, q: &DenseTensor, k: &DenseTensor, positions: Option<&[usize]>, ) -> (DenseTensor, DenseTensor)

Apply RoPE to query and key tensors

Arguments

• q - Query tensor [batch_size, seq_len, dim]
• k - Key tensor [batch_size, seq_len, dim]
• positions - Optional position indices

Returns

Tuple of (rotated_q, rotated_k)

pub fn dim(&self) -> usize

Get the dimension

pub fn max_seq_len(&self) -> usize

Get the maximum sequence length

Trait Implementations

impl Clone for RoPE

fn clone(&self) -> RoPE

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for RoPE

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

Formats the value using the given formatter.