Rotary Position Embedding (RoPE)

Background

Rotary Position Embedding (RoPE, Su et al. 2021) is the positional scheme that displaced sinusoidal and learned embeddings in modern LLMs — Llama, Mistral, Qwen, and Gemma all use it. Instead of adding a positional signal to the token embeddings, RoPE is applied directly to the query and key vectors before attention: it splits each feature vector into 2D pairs and rotates each pair by an angle that grows with position. Because a rotation preserves length, the attention dot product $q\cdot k$ ends up encoding both content similarity and relative position (through the angle difference).

Problem statement

Implement apply_rotary(x) for a (B, T, C) tensor (C even). Split the last axis into C/2 pairs; rotate the pair $(x_{2i}, x_{2i+1})$ at position $t$ by angle $\theta = t / 10000^{2i/C}$:

$$\theta = \frac{t}{10000^{\,2i/C}}, \qquad \begin{bmatrix} y_{2i} \\ y_{2i+1} \end{bmatrix} = \begin{bmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{bmatrix} \begin{bmatrix} x_{2i} \\ x_{2i+1} \end{bmatrix}$$

i.e. $y_{2i} = x_{2i}\cos\theta - x_{2i+1}\sin\theta$ and $y_{2i+1} = x_{2i}\sin\theta + x_{2i+1}\cos\theta$.

Input

• x — np.ndarray of shape (B, T, C); C must be even.

Output

Returns an np.ndarray of shape (B, T, C) — RoPE applied.

Examples

Example 1 — position 0 is the identity

Input:  any x; inspect position t=0
Output: out[:, 0, :] == x[:, 0, :]   (unchanged)

Explanation: at $t=0$ every angle $\theta = 0$, so $\cos\theta=1, \sin\theta=0$ and each pair is rotated by nothing — the first token passes through untouched.

Example 2 — hand-checked rotation ($C=4$, position $t=1$)

Input:  x[0, 1] = [1.0, 0.0, 1.0, 0.0]
Output: [cos(1), sin(1), cos(0.01), sin(0.01)] ≈ [0.5403, 0.8415, 0.99995, 0.01]

Explanation: with $C=4$ the two pair frequencies are $\theta_0 = t/10000^{0} = t$ and $\theta_1 = t/10000^{2/4} = t/100$. At $t=1$: pair 0 [1,0] rotates by 1 rad → [cos 1, sin 1]; pair 1 [1,0] rotates by 0.01 rad → [cos 0.01, sin 0.01].

Constraints

• C is even; pair the last axis as $(2i, 2i+1)$ and rotate each pair independently (no cross-pair mixing) — x.reshape(B, T, C//2, 2) is the clean way.
• Use the rotation angle $\theta = t / 10000^{2i/C}$ for pair $i$ at position $t$; broadcast cos/sin of the (T, C/2) angle table over the batch.
• Position 0 is the identity; the operation is norm-preserving — $\lVert\text{out}\rVert = \lVert x\rVert$ per position (atol≈1e-6).
• Must run cleanly (no NaN/Inf) at realistic sizes (e.g. C=32, T=16).

Notes

• Sign matters. Get the cos/sin signs wrong and the "rotation" becomes a reflection — norms are still preserved, but the position-0-identity check fails, which is exactly what the diagnostic catches.
• Frequency spread. Low-index pairs rotate fast ($\theta \propto t$), high-index pairs slowly ($\theta \propto t/10000$), so attention can read short- or long-range patterns from different frequency bands. Contrast with additive sinusoidal PE.