pub fn build_axial_rope_tables(
ids: &[(u32, u32)],
inv_freq: &[f32],
head_dim: usize,
) -> (Vec<f32>, Vec<f32>)Expand description
(cos, sin) tables of [num_tokens, head_dim] f32, in the layout the
rotate-half kernel reads.
inv_freq has spatial_dim/2 = head_dim/4 entries. The per-token frequency
row is cat([h*inv_freq, w*inv_freq]) (length head_dim/2) duplicated once
to fill head_dim — the reference’s cat([freq_hw, freq_hw], -1). There is
no partial-rotary split: every dimension of the 64-wide head is rotated.