model/rope.py

import torch


def precompute_rope_freqs(head_dim: int, max_seq_len: int, theta: float = 10_000.0):
    """Precompute the rotation angles used by RoPE (Section 4.5).

    Returns a complex tensor of shape (max_seq_len, head_dim // 2) where each
    entry encodes the rotation to apply at that position/frequency pair.
    """
    assert head_dim % 2 == 0, "RoPE requires an even head_dim"
    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
    positions = torch.arange(max_seq_len).float()
    angles = torch.outer(positions, freqs)  # (seq_len, head_dim/2)
    return torch.polar(torch.ones_like(angles), angles)  # complex64


def apply_rope(x: torch.Tensor, rope_freqs: torch.Tensor) -> torch.Tensor:
    """Apply rotary position embedding to a tensor of shape (B, n_heads, T, head_dim).

    rope_freqs should be pre-sliced to the current sequence length T before
    being passed in, i.e. rope_freqs[:T].
    """
    B, H, T, D = x.shape
    x_complex = torch.view_as_complex(x.float().reshape(B, H, T, D // 2, 2))
    freqs = rope_freqs.view(1, 1, T, D // 2)
    x_rotated = x_complex * freqs
    out = torch.view_as_real(x_rotated).reshape(B, H, T, D)
    return out.type_as(x)