Struct RotaryEmbedding 

pub struct RotaryEmbedding { /* private fields */ }

Rotary Positional Embedding layer

Precomputes cos/sin embeddings for all positions up to max_position_embeddings

Implementations

impl RotaryEmbedding

pub fn new( dim: usize, max_position_embeddings: usize, base: f32, device: &Device, ) -> Result<Self>

Create new RoPE embeddings

Arguments

• dim - Dimension of the embeddings (must be even)
• max_position_embeddings - Maximum sequence length
• base - Base for inverse frequencies (typically 10000.0)
• device - Device to create tensors on

pub fn forward(&self) -> Result<(Tensor, Tensor)>

Get the cached cos/sin embeddings

Returns

Tuple of (cos, sin) tensors with shape [max_position_embeddings, dim]

pub fn forward_with_len(&self, seq_len: usize) -> Result<(Tensor, Tensor)>

Get cos/sin embeddings for a specific sequence length

Arguments

• seq_len - Length of the sequence (must be <= max_position_embeddings)

Returns

Tuple of (cos, sin) tensors with shape [seq_len, dim]