Skip to main content

rotate_chunk

Function rotate_chunk 

Source
pub(crate) fn rotate_chunk<B: Backend>(
    v_bhlpr: F<B, 5>,
    rot: &RotPrim<B>,
    start: usize,
    transpose: bool,
) -> F<B, 5>
Expand description

Rotate the state_rank axis of one chunk of states by the per-token cumulative rotation. transpose = false is the de-rotation to the physical frame (R(−θ) / left-multiply by Q); transpose = true is its transpose (R(+θ) / left-multiply by conj(Q)), the map cotangents take back to the cache frame.

§Shapes

  • v_bhlpr: [batch, nheads, len, per_head_dim, state_rank]