Paper → implementation
See this paper in code.
Community-written snippets that make the paper’s ideas concrete. These are not official implementations from the paper’s authors.
Give a vector position without changing its length
Attention compares queries and keys with dot products. A dot product only sees the two vectors’ contents, so on its own attention cannot tell which token came first or how far apart two tokens are. RoPE rotates vector pairs by angles determined by token position. Input and output are [B, H, T, D]: batch, attention heads, tokens, and features per head. D must be even. positions [T] supplies one position per token, shared across this batch.
7 raise ValueError("RoPE needs an even head dimension")8 half = x.shape[-1] // 29 frequencies = base ** (-torch.arange(half, device=x.device).float() / half)10 angles = positions.to(device=x.device, dtype=torch.float32)[:, None] * frequencies11 cos = angles.cos().to(x.dtype)[None, None, :, :]12 sin = angles.sin().to(x.dtype)[None, None, :, :]13 first, second = x[..., :half], x[..., half:]Put RoPE inside a complete attention layer
Now begin with token features x [B, T, E]. Learned projections create queries, keys, and values; split E across H heads so D=E/H. Rotate queries and keys before computing their scores. This runnable module returns contextual token features [B, T, E]. It is one attention layer, not a complete language model.
16 17class RotaryAttention(nn.Module):18 def __init__(self, dim, heads):19 super().__init__()20 if dim % heads or (dim // heads) % 2:21 raise ValueError("Each head needs an even dimension")22 self.heads = headsYour place is saved automatically. Reaching the end marks this paper as read.
Page 1