Skip to main content

apply_interleaved_rope

Function apply_interleaved_rope 

Source
pub fn apply_interleaved_rope(
    x_data: &mut [f32],
    time_ids: &[usize],
    n_heads: usize,
    seq_len: usize,
    head_dim: usize,
    rope_dim: usize,
    cos_table: &[f32],
    sin_table: &[f32],
) -> Result<()>
Expand description

Apply interleaved (paired) rotary embeddings in place on host data.

Rotates pairs (2i, 2i+1) of the first rope_dim dims of each [n_heads, seq_len, head_dim] row using flat [max_pos * half_rope] tables. Shared by Moirai-2, TabICL, TabPFN, Lag-Llama (rank arg covers their dim ordering differences).