Skip to main content

rotate_state_rank_blocks

Function rotate_state_rank_blocks 

Source
pub fn rotate_state_rank_blocks<const D: usize, const DB: usize>(
    v: Tensor<D>,
    q: Tensor<DB>,
) -> Tensor<D>
Expand description

Apply a per-block quaternion rotation to the state_rank axis of v.

v has shape [..., state_rank] with state_rank = 4·J, viewed as J independent quaternion blocks; q has shape [..., J, 4] (one unit quaternion per block, same leading dims as v). Returns q ⊗ v per block, i.e. the rotation L_q applied within each 4-block, reshaped back to [..., state_rank].

This is the generalisation of RoPE’s per-pair 2×2 rotation to per-block 4×4. To rotate by the inverse cumulative rotation when absorbing into B/C (B̄ = Pᵀ B), pass q = conj(Qcum): rotate_state_rank_blocks(b, conj(qcum)).

DB must equal D + 1 (the block-split inserts the J axis).