qualia_core_db/solvers/
rope.rs

1//! Rotary Position Embedding (RoPE) — the STEM definition: a 2-D rotation.
2//!
3//! RoPE encodes token position by **rotating** each adjacent dimension pair `(2i, 2i+1)` of a
4//! head by a position- and frequency-dependent angle
5//!
6//! ```text
7//! θ_i = (pos / scale) · base^(−2i / head_dim)
8//! (x0, x1) ↦ (x0·cosθ − x1·sinθ,  x0·sinθ + x1·cosθ)
9//! ```
10//!
11//! That is exactly a rotation in the `(2i, 2i+1)` plane — the same rotation a geometric-algebra
12//! rotor (`super::geometric_algebra`) performs — applied per head. It is orthogonal, so it
13//! preserves the norm of every pair. Nothing proprietary: it is trigonometry. The LLM runtime's
14//! `rope_inplace` is an `f32` backend computing this same function.
15//!
16//! In place on a caller-owned slice; zero allocation.
17
18/// Apply interleaved ("normal"/llama) RoPE to `vec`, treated as `n_heads` consecutive blocks of
19/// `head_dim` elements. Each block's adjacent pairs `(2i, 2i+1)` are rotated by `θ_i`. `pos` is
20/// the token position, `base` the RoPE frequency base (e.g. 10000), `scale` the position scaling
21/// (≤ 0 or non-finite is treated as 1).
22pub fn rope_interleaved(
23    vec: &mut [f64],
24    n_heads: usize,
25    head_dim: usize,
26    pos: f64,
27    base: f64,
28    scale: f64,
29) {
30    let half = head_dim / 2;
31    if half == 0 {
32        return;
33    }
34    let scale = if scale > 0.0 && scale.is_finite() {
35        scale
36    } else {
37        1.0
38    };
39    let scaled_pos = pos / scale;
40    for head in 0..n_heads {
41        let off = head * head_dim;
42        if off + head_dim > vec.len() {
43            return;
44        }
45        for i in 0..half {
46            let theta = scaled_pos * base.powf(-2.0 * i as f64 / head_dim as f64);
47            let (s, c) = theta.sin_cos();
48            let x0 = vec[off + 2 * i];
49            let x1 = vec[off + 2 * i + 1];
50            vec[off + 2 * i] = x0 * c - x1 * s;
51            vec[off + 2 * i + 1] = x0 * s + x1 * c;
52        }
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    use super::*;
59
60    #[test]
61    fn position_zero_is_identity() {
62        // θ = 0 ⇒ no rotation.
63        let orig = [1.0, 2.0, 3.0, 4.0];
64        let mut v = orig;
65        rope_interleaved(&mut v, 1, 4, 0.0, 10000.0, 1.0);
66        for i in 0..4 {
67            assert!((v[i] - orig[i]).abs() < 1e-12);
68        }
69    }
70
71    #[test]
72    fn rotation_preserves_pair_norm() {
73        // Rotations are orthogonal: each (2i, 2i+1) pair keeps its magnitude.
74        let orig = [0.3, -0.7, 1.1, 2.0, -1.5, 0.4];
75        let mut v = orig;
76        rope_interleaved(&mut v, 1, 6, 5.0, 10000.0, 1.0);
77        for i in 0..3 {
78            let n0 = orig[2 * i] * orig[2 * i] + orig[2 * i + 1] * orig[2 * i + 1];
79            let n1 = v[2 * i] * v[2 * i] + v[2 * i + 1] * v[2 * i + 1];
80            assert!(
81                (n0 - n1).abs() < 1e-9,
82                "pair {i} norm changed: {n0} -> {n1}"
83            );
84        }
85    }
86
87    #[test]
88    fn quarter_turn_known_rotation() {
89        // For i = 0, θ_0 = scaled_pos. Choose scaled_pos = π/2 ⇒ (x0, x1) -> (-x1, x0).
90        let mut v = [2.0, 5.0];
91        rope_interleaved(&mut v, 1, 2, std::f64::consts::FRAC_PI_2, 10000.0, 1.0);
92        assert!((v[0] - (-5.0)).abs() < 1e-9, "x0 = {}", v[0]);
93        assert!((v[1] - 2.0).abs() < 1e-9, "x1 = {}", v[1]);
94    }
95
96    #[test]
97    fn per_head_blocks_independent() {
98        // Two heads of dim 2; head 0 and head 1 each rotate their own pair by the same angle.
99        let mut v = [1.0, 0.0, 0.0, 1.0];
100        rope_interleaved(&mut v, 2, 2, std::f64::consts::FRAC_PI_2, 10000.0, 1.0);
101        // head 0: (1,0) -> (0,1);  head 1: (0,1) -> (-1,0).
102        assert!((v[0] - 0.0).abs() < 1e-9 && (v[1] - 1.0).abs() < 1e-9);
103        assert!((v[2] - (-1.0)).abs() < 1e-9 && (v[3] - 0.0).abs() < 1e-9);
104    }
105}
qualia_core_db/solvers/rope.rs

qualia_core_db/solvers/
rope.rs
