feat: add rope visualization only with some minor cmparison with sinusoidial positional embedding

README.md

@@ -1,13 +1,50 @@
----
-title: ROPE Embedding visualization
-emoji: 📚
-colorFrom: gray
-colorTo: green
-sdk: gradio
-sdk_version: 6.26.0
-python_version: 3.11
-app_file: app.py
-pinned: false
-short_description: CVisualize how ROPE changes the embeddings
----
-# rope-implementation
+---
+title: RoPE Embedding Visualization
+emoji: 📚
+colorFrom: gray
+colorTo: green
+sdk: gradio
+sdk_version: 6.26.0
+python_version: 3.11
+app_file: app.py
+pinned: false
+short_description: Visualize how RoPE rotates query and key vectors
+---
+
+# RoPE Explorer
+
+Interactive Gradio app for **Rotary Position Embedding**. The Hugging Face Space
+runs `app.py` (`sdk: gradio`). Local Docker (`Dockerfile`, `compose.yaml`) is for
+running `python app.py` on port **7860**.
+
+Production math lives in [`src/`](src/) (`rope.py`, `absolute_pe.py`, `extract.py`,
+`plots.py`). Scratch scripts under `rope_implementation/`,
+`absolute_sinusoidal_position_embedding/`, and `relative_pos_embedding/` are
+learning notes only and are **not** imported by the app.
+
+## Modes
+
+1. **Random matrix** — sample even-width Q (and K) tensors, apply numpy RoPE,
+   inspect heatmaps, pairwise 2D rotation, `QK^T`, and additive sinusoidal PE.
+2. **Real model** — lazy-load an ungated Llama-like checkpoint (default
+   `HuggingFaceTB/SmolLM2-135M`), take `embed_tokens`, first-layer `q_proj` /
+   `k_proj` (GQA-aware), and compare educational numpy RoPE (`llama` pairing)
+   to the model's `rotary_emb`. No Hugging Face token is required. Gated models
+   are not used.
+
+First load of a model downloads weights into the cache; later runs reuse the
+last loaded model in memory.
+
+CPU is enough for SmolLM2 and Qwen2.5-0.5B. TinyLlama is included for a larger
+example and may be slow on CPU. This Space does **not** require ZeroGPU
+(`@spaces.GPU` is unused).
+
+## Local run
+
+```bash
+pip install -r requirements.txt
+python app.py
+```
+
+Or `docker compose up`. Hugging Face Cloud uses the README YAML (`sdk: gradio`),
+not the Docker image, unless the Space SDK is switched to Docker.

app.py

@@ -1,20 +1,353 @@
-import gradio as gr 
-import spaces
-
-@spaces.GPU
-def gpu_test():
-    return "GPU available"
-
-
-with gr.Blocks(title="RoPE Explorer") as demo: 
-    with gr.Tabs(): 
-        with gr.Tab("Block 01"): 
-            gr.Markdown("## Block 01") 
-            gr.Markdown("Placeholder for the first visualization.") 
-        
-        with gr.Tab("Block 02"): 
-            gr.Markdown("## Block 02") 
-            gr.Markdown("Placeholder for the second visualization.") 
-            
-if __name__ == "__main__": 
-    demo.launch( server_name="0.0.0.0", server_port=7860, )
+"""RoPE Explorer Gradio app. Imports only from ``src/``."""
+
+from __future__ import annotations
+
+import numpy as np
+import plotly.graph_objects as go
+import gradio as gr
+
+from src.absolute_pe import add_positional_encoding
+from src.extract import (
+    DEFAULT_MODEL,
+    MAX_SEQ_LEN,
+    MODEL_CHOICES,
+    expand_kv_heads,
+    extract_from_model,
+    random_qk,
+    select_head,
+)
+from src.plots import (
+    additive_pe_heatmaps,
+    attention_bars,
+    attention_heatmaps,
+    bulk_before_after_delta,
+    frequency_strip,
+    norm_compare_add_vs_rope,
+    norms_and_cosine,
+    position_sweep,
+    rotation_2d,
+    theta_heatmap,
+)
+from src.rope import (
+    attention_scores,
+    pair_dim_labels,
+    pair_frequencies,
+    pair_xy,
+    rotate_pair,
+    theta_grid,
+)
+
+HOWTO_MD = """
+# How RoPE works
+
+After each token has a vector from the **embedding table**, attention builds two extra
+vectors per token with linear layers (`q_proj`, `k_proj`):
+
+- **Q (query)** — what this token is looking for
+- **K (key)** — what this token offers as a match
+
+Attention scores are (scaled) **dot products** `Q · K`. **RoPE rotates those Q and K
+vectors in 2D planes before the dot product.** It does **not** add a position vector
+onto the raw token embeddings. This app shows embeddings only as context, then focuses
+on Q and K before vs after RoPE.
+
+## Pairwise rotation
+
+For even dimension `d`, pair `i` uses frequency
+
+$$\\omega_i = 10000^{-2i/d},\\qquad \\theta(k,i) = k\\,\\omega_i$$
+
+**Interleaved (paper-style)** pairing `(2i, 2i+1)`:
+
+$$
+x'_{2i} = x_{2i}\\cos\\theta - x_{2i+1}\\sin\\theta,\\qquad
+x'_{2i+1} = x_{2i}\\sin\\theta + x_{2i+1}\\cos\\theta
+$$
+
+Hugging Face Llama-like models use the same frequencies but pair `(i, i + d/2)`
+(`rotate_half`). Random-matrix mode uses interleaved pairing; real models use the
+Llama layout so the numpy implementation can be checksummed against `rotary_emb`.
+
+Relative positions fall out of the algebra: `R(m)^T R(n) = R(n-m)`.
+
+Shaw relative attention (learned bias `b_{m-n}` on scores) is a **different**
+mechanism and is not computed here.
+"""
+
+PLACEHOLDER = go.Figure().update_layout(
+    title="Run **Compute** on the Setup tab first",
+    template="plotly_white",
+    height=320,
+)
+
+
+def _safe_slider_max(n: int) -> int:
+    """Gradio sliders need max > min; keep a one-step range even at edge cases."""
+    return max(int(n), 1)
+
+
+def compute(
+    source: str,
+    sentence: str,
+    model_name: str,
+    seq_len: int,
+    dim: int,
+    seed: int,
+    base: float,
+    progress=gr.Progress(track_tqdm=False),
+):
+    try:
+        if source.startswith("Random"):
+            progress(0.4, desc="Sampling random Q/K")
+            data = random_qk(int(seq_len), int(dim), seed=int(seed), base=float(base))
+        else:
+            progress(0.2, desc=f"Loading {model_name} (first time downloads weights)")
+            data = extract_from_model(model_name, sentence)
+        seq = int(select_head(data["q_before"], 0).shape[0])
+        head_dim = int(select_head(data["q_before"], 0).shape[1])
+        n_pairs = head_dim // 2
+        n_heads = max(int(data["n_q_heads"]) - 1, 0)
+        checksum = data["checksum"]
+        if checksum is None:
+            status = (
+                f"Random Q/K · seq={seq} · dim={head_dim} · base={data['base']:g} · "
+                f"style={data['style']}"
+            )
+        else:
+            status = (
+                f"Model `{data['model_name']}` · {seq} tokens · head_dim={head_dim} · "
+                f"Q heads={data['n_q_heads']} · KV heads={data['n_kv_heads']} · "
+                f"rope_theta={data['base']:g} · "
+                f"max |numpy RoPE − model rotary| on Q = **{checksum:.3e}**"
+            )
+        token_labels = ", ".join(data["tokens"][:seq])
+        status = status + f"\n\nTokens: `{token_labels}`"
+        return (
+            data,
+            status,
+            gr.update(maximum=_safe_slider_max(n_heads), value=0),
+            gr.update(maximum=_safe_slider_max(seq - 1), value=0),
+            gr.update(maximum=_safe_slider_max(n_pairs - 1), value=0),
+            gr.update(maximum=_safe_slider_max(seq - 1), value=0),
+        )
+    except Exception as exc:
+        return (
+            None,
+            f"**Error:** {exc}",
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+        )
+
+
+def _qk_slice(data: dict, which: str, head: int):
+    before = data["q_before"] if which == "Q" else data["k_before"]
+    after = data["q_after"] if which == "Q" else data["k_after"]
+    return select_head(before, head), select_head(after, head)
+
+
+def update_bulk(data, which, head, mod_2pi):
+    if not data:
+        fig = PLACEHOLDER
+        return fig, fig, fig, fig
+    before, after = _qk_slice(data, which, int(head))
+    dim = before.shape[-1]
+    seq = before.shape[0]
+    return (
+        bulk_before_after_delta(before, after, tokens=data["tokens"]),
+        norms_and_cosine(before, after, tokens=data["tokens"]),
+        theta_heatmap(seq, dim, data["base"], mod_2pi=bool(mod_2pi)),
+        frequency_strip(dim, data["base"]),
+    )
+
+
+def update_individual(data, which, head, token, pair, sweep):
+    if not data:
+        return "Compute on the Setup tab first.", PLACEHOLDER, PLACEHOLDER
+    before, after = _qk_slice(data, which, int(head))
+    token = int(np.clip(token, 0, before.shape[0] - 1))
+    n_pairs = before.shape[1] // 2
+    pair = int(np.clip(pair, 0, n_pairs - 1))
+    style = data["style"]
+    xb, yb = pair_xy(before, token, pair, style=style)
+    xa, ya = pair_xy(after, token, pair, style=style)
+    theta = float(theta_grid(before.shape[0], before.shape[1], data["base"])[token, pair])
+    cos_t, sin_t = float(np.cos(theta)), float(np.sin(theta))
+    xe_chk, xo_chk = rotate_pair(np.array([xb]), np.array([yb]), np.array([theta]))
+    d0, d1 = pair_dim_labels(pair, before.shape[1], style=style)
+    table = f"""
+### Token `{token}` · pair `{pair}` (`{d0}`, `{d1}`)
+
+| | {d0} | {d1} |
+|---|---:|---:|
+| before | {xb:.6f} | {yb:.6f} |
+| after | {xa:.6f} | {ya:.6f} |
+| check (`rotate_pair`) | {float(xe_chk):.6f} | {float(xo_chk):.6f} |
+
+**θ(k,i) = {theta:.6f} rad** · cos = {cos_t:.6f} · sin = {sin_t:.6f}
+
+`x'_even = x_even cos θ − x_odd sin θ`  
+`x'_odd  = x_even sin θ + x_odd cos θ`
+"""
+    neighbors = [n for n in (token - 1, token + 1, token + 2) if 0 <= n < before.shape[0]]
+    rot = rotation_2d(before, after, token, pair, style, theta, neighbor_tokens=neighbors)
+    if sweep:
+        omega = float(pair_frequencies(before.shape[1], data["base"])[pair])
+        sweep_fig = position_sweep(xb, yb, omega, before.shape[0], token)
+    else:
+        sweep_fig = PLACEHOLDER
+        sweep_fig.update_layout(title="Enable “replay same pair at every k” to see position-only spin")
+    return table, rot, sweep_fig
+
+
+def update_attention(data, head, query_token):
+    if not data:
+        return PLACEHOLDER, PLACEHOLDER, ""
+    q_b = select_head(data["q_before"], int(head))
+    q_a = select_head(data["q_after"], int(head))
+    k_b_all = expand_kv_heads(data["k_before"], data["n_q_heads"])
+    k_a_all = expand_kv_heads(data["k_after"], data["n_q_heads"])
+    k_b = select_head(k_b_all, int(head))
+    k_a = select_head(k_a_all, int(head))
+    sb = attention_scores(q_b, k_b)
+    sa = attention_scores(q_a, k_a)
+    qt = int(np.clip(query_token, 0, sb.shape[0] - 1))
+    note = (
+        "Additive PE changes values by **addition**. RoPE encodes **relative** offset "
+        "because `R(m)^T R(n) = R(n−m)`: the score depends on the position difference, "
+        "not on absolute indices alone."
+    )
+    return attention_heatmaps(sb, sa), attention_bars(sb[qt], sa[qt], qt), note
+
+
+def update_compare(data):
+    if not data:
+        return PLACEHOLDER, PLACEHOLDER, ""
+    emb = np.asarray(data["embeddings"], dtype=np.float64)
+    # Compare tab always uses additive PE on the embedding matrix (may be wider than a head).
+    pe, combined = add_positional_encoding(emb, base=data["base"])
+    q_b = select_head(data["q_before"], 0)
+    q_a = select_head(data["q_after"], 0)
+    heat = additive_pe_heatmaps(emb, pe, combined)
+    norms = norm_compare_add_vs_rope(emb, combined, q_b, q_a)
+    copy = """
+**Absolute sinusoidal PE** *adds* a position-shaped vector, so both **norm and direction** change.
+
+**RoPE** *rotates* query/key pairs: **norm stays**, and the relative angle depends on `m − n`.
+
+Shaw-style relative attention (`q_m^T k_n + b_{m-n}`) is a third, learned-bias mechanism — not shown as a plot.
+"""
+    return heat, norms, copy
+
+
+def toggle_source(source: str):
+    is_random = source.startswith("Random")
+    return (
+        gr.update(visible=is_random),
+        gr.update(visible=is_random),
+        gr.update(visible=is_random),
+        gr.update(visible=not is_random),
+        gr.update(visible=not is_random),
+    )
+
+
+with gr.Blocks(title="RoPE Explorer") as demo:
+    state = gr.State(None)
+    gr.Markdown("# RoPE Explorer")
+    gr.Markdown(
+        "Interactive view of **Rotary Position Embedding**: random Q/K matrices or "
+        "query/key vectors from a small ungated Hugging Face model."
+    )
+    with gr.Tabs():
+        with gr.Tab("How RoPE works"):
+            gr.Markdown(HOWTO_MD)
+
+        with gr.Tab("Setup"):
+            source = gr.Radio(
+                ["Random matrix", "Real model"],
+                value="Random matrix",
+                label="Source",
+            )
+            with gr.Row():
+                sentence = gr.Textbox(
+                    value="RoPE rotates query and key vectors.",
+                    label="Sentence (real model)",
+                    visible=False,
+                )
+                model_name = gr.Dropdown(
+                    MODEL_CHOICES,
+                    value=DEFAULT_MODEL,
+                    label="Model (ungated, Llama-like)",
+                    visible=False,
+                )
+            with gr.Row():
+                seq_len = gr.Slider(2, MAX_SEQ_LEN, value=16, step=1, label="Sequence length")
+                dim = gr.Slider(4, 128, value=32, step=2, label="Dimension (even)")
+                seed = gr.Number(value=0, label="Seed", precision=0)
+            base = gr.Number(
+                value=10000,
+                label="RoPE base (overridden by config.rope_theta for real models)",
+            )
+            compute_btn = gr.Button("Compute", variant="primary")
+            status = gr.Markdown("Choose a source and click Compute.")
+            head = gr.Slider(minimum=0, maximum=2, step=1, value=0, label="Head index (real models)")
+
+        with gr.Tab("Bulk changes"):
+            which = gr.Radio(["Q", "K"], value="Q", label="Tensor")
+            mod_2pi = gr.Checkbox(False, label="θ heatmap: wrap mod 2π")
+            bulk_main = gr.Plot(label="Before / after / delta")
+            bulk_norm = gr.Plot(label="Norms and cosine")
+            bulk_theta = gr.Plot(label="θ(k, i)")
+            bulk_freq = gr.Plot(label="ω_i")
+
+        with gr.Tab("Individual changes"):
+            with gr.Row():
+                token_k = gr.Slider(0, 15, step=1, value=0, label="Token index k")
+                pair_i = gr.Slider(0, 15, step=1, value=0, label="Pair index i")
+            sweep = gr.Checkbox(True, label="Replay the same content pair at every position k")
+            pair_table = gr.Markdown()
+            pair_plot = gr.Plot()
+            sweep_plot = gr.Plot()
+
+        with gr.Tab("Attention effect"):
+            query_token = gr.Slider(0, 15, step=1, value=0, label="Query token")
+            attn_heat = gr.Plot()
+            attn_bar = gr.Plot()
+            attn_note = gr.Markdown()
+
+        with gr.Tab("Compare to additive PE"):
+            pe_heat = gr.Plot()
+            pe_norm = gr.Plot()
+            pe_note = gr.Markdown()
+
+    compute_btn.click(
+        compute,
+        inputs=[source, sentence, model_name, seq_len, dim, seed, base],
+        outputs=[state, status, head, token_k, pair_i, query_token],
+    )
+    source.change(
+        toggle_source,
+        inputs=[source],
+        outputs=[seq_len, dim, seed, sentence, model_name],
+    )
+
+    bulk_inputs = [state, which, head, mod_2pi]
+    bulk_outputs = [bulk_main, bulk_norm, bulk_theta, bulk_freq]
+    for ctrl in bulk_inputs:
+        ctrl.change(update_bulk, inputs=bulk_inputs, outputs=bulk_outputs)
+
+    ind_inputs = [state, which, head, token_k, pair_i, sweep]
+    ind_outputs = [pair_table, pair_plot, sweep_plot]
+    for ctrl in ind_inputs:
+        ctrl.change(update_individual, inputs=ind_inputs, outputs=ind_outputs)
+
+    attn_inputs = [state, head, query_token]
+    attn_outputs = [attn_heat, attn_bar, attn_note]
+    for ctrl in attn_inputs:
+        ctrl.change(update_attention, inputs=attn_inputs, outputs=attn_outputs)
+
+    state.change(update_compare, inputs=[state], outputs=[pe_heat, pe_norm, pe_note])
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)

requirements.txt

@@ -1,5 +1,7 @@
-gradio==6.26.0
-
-numpy>=2.0,<2.3
-spaces==0.51.3
-plotly==7.0.0
+gradio==6.26.0
+numpy>=2.0,<2.3
+plotly==7.0.0
+spaces==0.51.3
+transformers>=4.44.0
+accelerate>=0.34.0
+torch>=2.2.0

src/absolute_pe.py

@@ -0,0 +1,35 @@
+"""Vectorized additive sinusoidal positional encoding (Vaswani et al.).
+
+``p(k, i) = sin(k / base^{i/d})`` for even ``i``, and
+``p(k, i) = cos(k / base^{(i-1)/d})`` for odd ``i``.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def sinusoidal_pe(seq_len: int, dim: int, base: float = 10000.0) -> np.ndarray:
+    """Return a ``(seq_len, dim)`` additive PE matrix."""
+    if seq_len < 0 or dim < 1:
+        raise ValueError("seq_len must be >= 0 and dim >= 1")
+    positions = np.arange(seq_len, dtype=np.float64)[:, None]
+    dims = np.arange(dim, dtype=np.float64)[None, :]
+    exponent = np.where(dims % 2 == 0, dims / dim, (dims - 1.0) / dim)
+    angles = positions / (base ** exponent)
+    pe = np.empty((seq_len, dim), dtype=np.float64)
+    pe[:, 0::2] = np.sin(angles[:, 0::2])
+    pe[:, 1::2] = np.cos(angles[:, 1::2])
+    return pe
+
+
+def add_positional_encoding(
+    embeddings: np.ndarray,
+    base: float = 10000.0,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Return ``(pe, embeddings + pe)`` for a ``(seq, dim)`` matrix."""
+    embeddings = np.asarray(embeddings, dtype=np.float64)
+    if embeddings.ndim != 2:
+        raise ValueError("embeddings must be 2D (seq, dim)")
+    pe = sinusoidal_pe(embeddings.shape[0], embeddings.shape[1], base=base)
+    return pe, embeddings + pe

src/extract.py

@@ -0,0 +1,260 @@
+"""Random tensors and lazy Hugging Face Q/K extraction.
+
+Models are not loaded at import time. The last loaded model is cached.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from src.rope import apply_rope
+
+MAX_SEQ_LEN = 64
+
+# Ungated Llama-like checkpoints (q_proj / k_proj + rotary). No HF token required.
+MODEL_CHOICES = [
+    "HuggingFaceTB/SmolLM2-135M",
+    "Qwen/Qwen2.5-0.5B-Instruct",
+    "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+]
+DEFAULT_MODEL = MODEL_CHOICES[0]
+
+_cache: dict[str, Any] = {"name": None, "model": None, "tokenizer": None}
+
+
+def random_matrix(seq_len: int, dim: int, seed: int = 0) -> np.ndarray:
+    if dim % 2 != 0:
+        raise ValueError(f"dim must be even for RoPE, got {dim}")
+    seq_len = int(np.clip(seq_len, 1, MAX_SEQ_LEN))
+    rng = np.random.default_rng(int(seed))
+    return rng.standard_normal((seq_len, dim)).astype(np.float64)
+
+
+def random_qk(
+    seq_len: int,
+    dim: int,
+    seed: int = 0,
+    base: float = 10000.0,
+) -> dict[str, Any]:
+    q = random_matrix(seq_len, dim, seed=seed)
+    k = random_matrix(seq_len, dim, seed=seed + 1)
+    return _pack_tensors(
+        q_before=q,
+        k_before=k,
+        embeddings=q.copy(),
+        tokens=[f"t{i}" for i in range(q.shape[0])],
+        base=float(base),
+        style="interleaved",
+        n_q_heads=1,
+        n_kv_heads=1,
+        checksum=None,
+        source="random",
+        model_name=None,
+    )
+
+
+def _pack_tensors(
+    *,
+    q_before: np.ndarray,
+    k_before: np.ndarray,
+    embeddings: np.ndarray,
+    tokens: list[str],
+    base: float,
+    style: str,
+    n_q_heads: int,
+    n_kv_heads: int,
+    checksum: float | None,
+    source: str,
+    model_name: str | None,
+    q_after_model: np.ndarray | None = None,
+    k_after_model: np.ndarray | None = None,
+) -> dict[str, Any]:
+    q_after = apply_rope(q_before, base=base, style=style)
+    k_after = apply_rope(k_before, base=base, style=style)
+    return {
+        "q_before": q_before,
+        "k_before": k_before,
+        "q_after": q_after,
+        "k_after": k_after,
+        "q_after_model": q_after_model,
+        "k_after_model": k_after_model,
+        "embeddings": embeddings,
+        "tokens": tokens,
+        "base": float(base),
+        "style": style,
+        "n_q_heads": int(n_q_heads),
+        "n_kv_heads": int(n_kv_heads),
+        "checksum": checksum,
+        "source": source,
+        "model_name": model_name,
+    }
+
+
+def _require_hf():
+    try:
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+    except ImportError as exc:
+        raise RuntimeError(
+            "Real-model mode needs `torch` and `transformers`. "
+            "Install them or use Random matrix mode."
+        ) from exc
+    return torch, AutoModelForCausalLM, AutoTokenizer
+
+
+def get_model(model_name: str):
+    """Load tokenizer + causal LM on CPU; cache the last selection."""
+    torch, AutoModelForCausalLM, AutoTokenizer = _require_hf()
+    if _cache["name"] == model_name and _cache["model"] is not None:
+        return _cache["model"], _cache["tokenizer"]
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True,
+    )
+    model.eval()
+    model.to("cpu")
+    _cache["name"] = model_name
+    _cache["model"] = model
+    _cache["tokenizer"] = tokenizer
+    return model, tokenizer
+
+
+def _backbone(model):
+    if hasattr(model, "model") and hasattr(model.model, "embed_tokens"):
+        return model.model
+    raise RuntimeError(
+        "This checkpoint is not Llama-like (expected model.model.embed_tokens)."
+    )
+
+
+def _rotary_module(backbone, attn):
+    if hasattr(backbone, "rotary_emb"):
+        return backbone.rotary_emb
+    if hasattr(attn, "rotary_emb"):
+        return attn.rotary_emb
+    return None
+
+
+def _rotate_half_torch(x, torch):
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _broadcast_cos_sin(cos, sin, q):
+    """Make cos/sin broadcast with q of shape (batch, heads, seq, dim)."""
+    if cos.dim() == 2:
+        cos, sin = cos.unsqueeze(0).unsqueeze(0), sin.unsqueeze(0).unsqueeze(0)
+    elif cos.dim() == 3:
+        cos, sin = cos.unsqueeze(1), sin.unsqueeze(1)
+    return cos, sin
+
+
+def _hf_rotary(q, k, rotary, position_ids, torch):
+    try:
+        cos, sin = rotary(q, position_ids=position_ids)
+    except TypeError:
+        try:
+            cos, sin = rotary(q, seq_len=q.shape[-2])
+        except TypeError:
+            cos, sin = rotary(q)
+    if cos.shape[-1] == q.shape[-1] // 2:
+        cos = torch.cat((cos, cos), dim=-1)
+        sin = torch.cat((sin, sin), dim=-1)
+    cos, sin = _broadcast_cos_sin(cos, sin, q)
+    q_rot = (q * cos) + (_rotate_half_torch(q, torch) * sin)
+    k_rot = (k * cos) + (_rotate_half_torch(k, torch) * sin)
+    return q_rot, k_rot
+
+
+def extract_from_model(model_name: str, sentence: str) -> dict[str, Any]:
+    torch, _, _ = _require_hf()
+    model, tokenizer = get_model(model_name)
+    text = sentence.strip() or "RoPE rotates query and key vectors."
+    encoded = tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        max_length=MAX_SEQ_LEN,
+        add_special_tokens=True,
+    )
+    input_ids = encoded["input_ids"]
+    tokens = tokenizer.convert_ids_to_tokens(input_ids[0].tolist())
+    backbone = _backbone(model)
+    cfg = model.config
+    n_heads = int(cfg.num_attention_heads)
+    n_kv = int(getattr(cfg, "num_key_value_heads", n_heads))
+    hidden_size = int(cfg.hidden_size)
+    head_dim = int(getattr(cfg, "head_dim", hidden_size // n_heads))
+    base = float(getattr(cfg, "rope_theta", 10000.0))
+
+    with torch.no_grad():
+        embeds = backbone.embed_tokens(input_ids)
+        layer = backbone.layers[0]
+        attn = layer.self_attn
+        hidden = layer.input_layernorm(embeds)
+        q = attn.q_proj(hidden)
+        k = attn.k_proj(hidden)
+        seq = q.shape[1]
+        q = q.view(1, seq, n_heads, head_dim).transpose(1, 2).contiguous()
+        k = k.view(1, seq, n_kv, head_dim).transpose(1, 2).contiguous()
+        position_ids = torch.arange(seq).unsqueeze(0)
+        rotary = _rotary_module(backbone, attn)
+        q_model = k_model = None
+        if rotary is not None:
+            q_model, k_model = _hf_rotary(q, k, rotary, position_ids, torch)
+
+    q_np = q[0].cpu().numpy().astype(np.float64)
+    k_np = k[0].cpu().numpy().astype(np.float64)
+    emb_np = embeds[0].cpu().numpy().astype(np.float64)
+    q_model_np = k_model_np = None
+    checksum = None
+    if q_model is not None:
+        q_model_np = q_model[0].cpu().numpy().astype(np.float64)
+        k_model_np = k_model[0].cpu().numpy().astype(np.float64)
+        q_edu = apply_rope(q_np, base=base, style="llama")
+        checksum = float(np.max(np.abs(q_edu - q_model_np)))
+
+    return _pack_tensors(
+        q_before=q_np,
+        k_before=k_np,
+        embeddings=emb_np,
+        tokens=tokens,
+        base=base,
+        style="llama",
+        n_q_heads=n_heads,
+        n_kv_heads=n_kv,
+        checksum=checksum,
+        source="model",
+        model_name=model_name,
+        q_after_model=q_model_np,
+        k_after_model=k_model_np,
+    )
+
+
+def select_head(tensor: np.ndarray, head: int) -> np.ndarray:
+    """Return a ``(seq, dim)`` slice; ``tensor`` is 2D or ``(heads, seq, dim)``."""
+    t = np.asarray(tensor)
+    if t.ndim == 2:
+        return t
+    if t.ndim != 3:
+        raise ValueError(f"expected 2D or 3D tensor, got {t.ndim}D")
+    h = int(np.clip(head, 0, t.shape[0] - 1))
+    return t[h]
+
+
+def expand_kv_heads(k: np.ndarray, n_q_heads: int) -> np.ndarray:
+    """Repeat GQA key heads so they align with query heads."""
+    t = np.asarray(k)
+    if t.ndim == 2:
+        return t
+    n_kv = t.shape[0]
+    if n_kv == n_q_heads:
+        return t
+    if n_q_heads % n_kv != 0:
+        return t
+    return np.repeat(t, n_q_heads // n_kv, axis=0)

src/plots.py

@@ -0,0 +1,305 @@
+"""Plotly figures for bulk heatmaps, pair rotation, and attention."""
+
+from __future__ import annotations
+
+import numpy as np
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+
+from src.rope import l2_norms, pair_frequencies, pair_xy, row_cosine, theta_grid
+
+
+def _empty(title: str) -> go.Figure:
+    fig = go.Figure()
+    fig.update_layout(title=title, template="plotly_white")
+    return fig
+
+
+def downsample(mat: np.ndarray, max_cols: int = 64, max_rows: int = 64) -> np.ndarray:
+    m = np.asarray(mat)
+    row_step = max(1, int(np.ceil(m.shape[0] / max_rows)))
+    if m.ndim == 1:
+        return m[::row_step]
+    col_step = max(1, int(np.ceil(m.shape[1] / max_cols)))
+    return m[::row_step, ::col_step]
+
+
+def heatmap(z: np.ndarray, title: str, xtitle: str, ytitle: str, tokens=None) -> go.Figure:
+    z = downsample(np.asarray(z, dtype=np.float64))
+    fig = go.Figure(
+        data=go.Heatmap(
+            z=z,
+            colorbar=dict(title="value"),
+            hovertemplate="y=%{y}<br>x=%{x}<br>z=%{z:.4f}<extra></extra>",
+        )
+    )
+    fig.update_layout(
+        title=title,
+        xaxis_title=xtitle,
+        yaxis_title=ytitle,
+        yaxis=dict(autorange="reversed"),
+        template="plotly_white",
+        height=380,
+        margin=dict(l=60, r=40, t=50, b=50),
+    )
+    if tokens is not None and len(tokens) == z.shape[0]:
+        fig.update_yaxes(tickmode="array", tickvals=list(range(len(tokens))), ticktext=tokens)
+    return fig
+
+
+def bulk_before_after_delta(before: np.ndarray, after: np.ndarray, tokens=None) -> go.Figure:
+    before = np.asarray(before, dtype=np.float64)
+    after = np.asarray(after, dtype=np.float64)
+    delta = after - before
+    mats = [downsample(before), downsample(after), downsample(delta)]
+    titles = ["Before RoPE", "After RoPE", "Delta (after − before)"]
+    fig = make_subplots(rows=1, cols=3, subplot_titles=titles)
+    for i, mat in enumerate(mats, start=1):
+        fig.add_trace(
+            go.Heatmap(
+                z=mat,
+                showscale=(i == 3),
+                hovertemplate="token=%{y}<br>dim=%{x}<br>z=%{z:.4f}<extra></extra>",
+            ),
+            row=1,
+            col=i,
+        )
+        fig.update_yaxes(autorange="reversed", row=1, col=i)
+    fig.update_layout(template="plotly_white", height=400, margin=dict(t=60))
+    return fig
+
+
+def norms_and_cosine(before: np.ndarray, after: np.ndarray, tokens=None) -> go.Figure:
+    nb = l2_norms(before)
+    na = l2_norms(after)
+    cos = row_cosine(before, after)
+    xs = list(range(len(nb)))
+    fig = make_subplots(rows=1, cols=2, subplot_titles=["Per-token L2 norm", "Cosine (original vs rotated)"])
+    fig.add_trace(go.Scatter(x=xs, y=nb, name="before", mode="lines+markers"), row=1, col=1)
+    fig.add_trace(go.Scatter(x=xs, y=na, name="after", mode="lines+markers"), row=1, col=1)
+    fig.add_trace(go.Bar(x=xs, y=cos, name="cosine", showlegend=False), row=1, col=2)
+    fig.update_xaxes(title_text="token", row=1, col=1)
+    fig.update_xaxes(title_text="token", row=1, col=2)
+    fig.update_yaxes(title_text="L2", row=1, col=1)
+    fig.update_yaxes(title_text="cosine", range=[min(0.0, float(np.min(cos)) - 0.05), 1.02], row=1, col=2)
+    fig.update_layout(template="plotly_white", height=380, barmode="group")
+    return fig
+
+
+def theta_heatmap(seq_len: int, dim: int, base: float, mod_2pi: bool = False) -> go.Figure:
+    grid = theta_grid(seq_len, dim, base=base)
+    if mod_2pi:
+        grid = np.mod(grid, 2 * np.pi)
+        title = "θ(k, i) mod 2π"
+    else:
+        title = "θ(k, i) = k · ω_i"
+    z = downsample(grid, max_cols=64, max_rows=64)
+    fig = go.Figure(
+        data=go.Heatmap(
+            z=z,
+            colorbar=dict(title="radians"),
+            hovertemplate="token k=%{y}<br>pair i=%{x}<br>θ=%{z:.4f}<extra></extra>",
+        )
+    )
+    fig.update_layout(
+        title=title,
+        xaxis_title="pair index i",
+        yaxis_title="token position k",
+        yaxis=dict(autorange="reversed"),
+        template="plotly_white",
+        height=380,
+    )
+    return fig
+
+
+def frequency_strip(dim: int, base: float) -> go.Figure:
+    omega = pair_frequencies(dim, base=base)
+    fig = go.Figure(data=go.Bar(x=list(range(len(omega))), y=omega, name="ω_i"))
+    fig.update_layout(
+        title="Pair frequencies ω_i (pair 0 is fastest)",
+        xaxis_title="pair index i",
+        yaxis_title="ω",
+        yaxis_type="log",
+        template="plotly_white",
+        height=280,
+    )
+    return fig
+
+
+def _arc_points(x0, y0, x1, y1, n=40):
+    r0 = float(np.hypot(x0, y0))
+    r1 = float(np.hypot(x1, y1))
+    r = 0.5 * (r0 + r1)
+    if r < 1e-9:
+        return [], []
+    a0 = float(np.arctan2(y0, x0))
+    a1 = float(np.arctan2(y1, x1))
+    delta = (a1 - a0 + np.pi) % (2 * np.pi) - np.pi
+    ts = np.linspace(0.0, 1.0, n)
+    angs = a0 + ts * delta
+    rr = r * 0.55
+    return list(rr * np.cos(angs)), list(rr * np.sin(angs))
+
+
+def rotation_2d(
+    matrix_before: np.ndarray,
+    matrix_after: np.ndarray,
+    token: int,
+    pair: int,
+    style: str,
+    theta: float,
+    neighbor_tokens: list[int] | None = None,
+) -> go.Figure:
+    xb, yb = pair_xy(matrix_before, token, pair, style=style)
+    xa, ya = pair_xy(matrix_after, token, pair, style=style)
+    fig = go.Figure()
+    fig.add_trace(
+        go.Scatter(
+            x=[0, xb],
+            y=[0, yb],
+            mode="lines+markers",
+            name="before",
+            line=dict(width=3),
+        )
+    )
+    fig.add_trace(
+        go.Scatter(
+            x=[0, xa],
+            y=[0, ya],
+            mode="lines+markers",
+            name="after",
+            line=dict(width=3),
+        )
+    )
+    xs, ys = _arc_points(xb, yb, xa, ya)
+    if xs:
+        fig.add_trace(
+            go.Scatter(x=xs, y=ys, mode="lines", name=f"θ={theta:.3f} rad", line=dict(dash="dot"))
+        )
+    if neighbor_tokens:
+        for n in neighbor_tokens:
+            if n == token:
+                continue
+            nx, ny = pair_xy(matrix_after, n, pair, style=style)
+            fig.add_trace(
+                go.Scatter(
+                    x=[0, nx],
+                    y=[0, ny],
+                    mode="lines+markers",
+                    name=f"token {n} after",
+                    opacity=0.45,
+                )
+            )
+    lim = max(abs(xb), abs(yb), abs(xa), abs(ya), 1e-3) * 1.3
+    fig.update_layout(
+        title=f"Pair {pair} at token {token}: 2D rotation",
+        xaxis=dict(scaleanchor="y", scaleratio=1, range=[-lim, lim], zeroline=True),
+        yaxis=dict(range=[-lim, lim], zeroline=True),
+        template="plotly_white",
+        height=460,
+        legend=dict(orientation="h"),
+    )
+    return fig
+
+
+def position_sweep(
+    x_even: float,
+    x_odd: float,
+    omega: float,
+    seq_len: int,
+    highlight: int,
+) -> go.Figure:
+    """Same content pair spun by θ(k)=k ω — position is the only change."""
+    ks = np.arange(seq_len)
+    thetas = ks * omega
+    xs = x_even * np.cos(thetas) - x_odd * np.sin(thetas)
+    ys = x_even * np.sin(thetas) + x_odd * np.cos(thetas)
+    fig = go.Figure()
+    fig.add_trace(go.Scatter(x=xs, y=ys, mode="markers+lines", name="path over k"))
+    fig.add_trace(
+        go.Scatter(
+            x=[0, xs[highlight]],
+            y=[0, ys[highlight]],
+            mode="lines+markers",
+            name=f"k={highlight}",
+            line=dict(width=3),
+        )
+    )
+    lim = max(float(np.max(np.abs(xs))), float(np.max(np.abs(ys))), 1e-3) * 1.3
+    fig.update_layout(
+        title="Same (x_even, x_odd) rotated at every position k",
+        xaxis=dict(scaleanchor="y", scaleratio=1, range=[-lim, lim]),
+        yaxis=dict(range=[-lim, lim]),
+        template="plotly_white",
+        height=420,
+    )
+    return fig
+
+
+def attention_heatmaps(scores_before: np.ndarray, scores_after: np.ndarray) -> go.Figure:
+    fig = make_subplots(rows=1, cols=2, subplot_titles=["QKᵀ without RoPE", "QKᵀ with RoPE"])
+    for i, mat in enumerate([scores_before, scores_after], start=1):
+        fig.add_trace(
+            go.Heatmap(
+                z=downsample(mat),
+                showscale=(i == 2),
+                hovertemplate="query=%{y}<br>key=%{x}<br>score=%{z:.4f}<extra></extra>",
+            ),
+            row=1,
+            col=i,
+        )
+        fig.update_yaxes(autorange="reversed", title_text="query token", row=1, col=i)
+        fig.update_xaxes(title_text="key token", row=1, col=i)
+    fig.update_layout(template="plotly_white", height=420)
+    return fig
+
+
+def attention_bars(logits_before: np.ndarray, logits_after: np.ndarray, query_token: int) -> go.Figure:
+    xs = list(range(len(logits_before)))
+    fig = go.Figure()
+    fig.add_trace(go.Bar(x=xs, y=logits_before, name="without RoPE"))
+    fig.add_trace(go.Bar(x=xs, y=logits_after, name="with RoPE"))
+    fig.update_layout(
+        title=f"Attention logits from query token {query_token}",
+        xaxis_title="key token",
+        yaxis_title="Q·K",
+        barmode="group",
+        template="plotly_white",
+        height=360,
+    )
+    return fig
+
+
+def additive_pe_heatmaps(emb: np.ndarray, pe: np.ndarray, combined: np.ndarray) -> go.Figure:
+    titles = ["Token embeddings", "Additive sinusoidal PE", "Embeddings + PE"]
+    fig = make_subplots(rows=1, cols=3, subplot_titles=titles)
+    for i, mat in enumerate([emb, pe, combined], start=1):
+        fig.add_trace(
+            go.Heatmap(z=downsample(mat), showscale=(i == 3)),
+            row=1,
+            col=i,
+        )
+        fig.update_yaxes(autorange="reversed", row=1, col=i)
+    fig.update_layout(template="plotly_white", height=400)
+    return fig
+
+
+def norm_compare_add_vs_rope(
+    emb: np.ndarray,
+    emb_plus_pe: np.ndarray,
+    q_before: np.ndarray,
+    q_after: np.ndarray,
+) -> go.Figure:
+    fig = go.Figure()
+    fig.add_trace(go.Scatter(y=l2_norms(emb), mode="lines+markers", name="embeddings"))
+    fig.add_trace(go.Scatter(y=l2_norms(emb_plus_pe), mode="lines+markers", name="embeddings + PE"))
+    fig.add_trace(go.Scatter(y=l2_norms(q_before), mode="lines+markers", name="Q before RoPE"))
+    fig.add_trace(go.Scatter(y=l2_norms(q_after), mode="lines+markers", name="Q after RoPE"))
+    fig.update_layout(
+        title="L2 norms: additive PE changes magnitude; RoPE does not",
+        xaxis_title="token",
+        yaxis_title="L2",
+        template="plotly_white",
+        height=360,
+    )
+    return fig

src/rope.py

@@ -0,0 +1,124 @@
+"""Vectorized RoPE: pairwise 2D rotations of query/key slices.
+
+Two pairing conventions:
+
+- ``interleaved`` — paper / GPT-J style: rotate dims ``(2i, 2i+1)``.
+- ``llama`` — Hugging Face Llama/Qwen/SmolLM style: rotate
+  ``(i, i + dim/2)`` (the ``rotate_half`` layout).
+
+Frequencies match the usual schedule ``ω_i = base^{-2i/d}``.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+VALID_STYLES = ("interleaved", "llama")
+
+
+def pair_frequencies(dim: int, base: float = 10000.0) -> np.ndarray:
+    """Return ``ω`` of shape ``(dim/2,)`` with ``ω_i = base ** (-2i / dim)``."""
+    if dim % 2 != 0:
+        raise ValueError(f"RoPE requires an even dim, got {dim}")
+    n_pairs = dim // 2
+    pair_idx = np.arange(n_pairs, dtype=np.float64) * 2.0
+    return 1.0 / (base ** (pair_idx / dim))
+
+
+def theta_grid(seq_len: int, dim: int, base: float = 10000.0) -> np.ndarray:
+    """Rotation angles ``θ[k, i] = k * ω_i``, shape ``(seq_len, dim/2)``."""
+    omega = pair_frequencies(dim, base=base)
+    positions = np.arange(seq_len, dtype=np.float64)[:, None]
+    return positions * omega[None, :]
+
+
+def _broadcast_angles(angles: np.ndarray, like: np.ndarray) -> np.ndarray:
+    """Broadcast ``(seq, n_pairs)`` onto ``like`` with shape ``(..., seq, n_pairs)``."""
+    lead = like.ndim - 2
+    return angles.reshape((1,) * lead + angles.shape)
+
+
+def rotate_pair(x_even: np.ndarray, x_odd: np.ndarray, theta: np.ndarray):
+    """Apply a 2D rotation by ``theta`` (radians) to coordinate pairs."""
+    cos = np.cos(theta)
+    sin = np.sin(theta)
+    return x_even * cos - x_odd * sin, x_even * sin + x_odd * cos
+
+
+def apply_rope(
+    x: np.ndarray,
+    base: float = 10000.0,
+    style: str = "interleaved",
+) -> np.ndarray:
+    """Rotate last-dimension pairs of ``x`` (``(..., seq_len, dim)``)."""
+    if style not in VALID_STYLES:
+        raise ValueError(f"style must be one of {VALID_STYLES}, got {style!r}")
+    x = np.asarray(x, dtype=np.float64)
+    if x.ndim < 2:
+        raise ValueError("x must be at least 2D (seq, dim)")
+    dim = x.shape[-1]
+    seq_len = x.shape[-2]
+    angles = theta_grid(seq_len, dim, base=base)
+    cos = np.cos(angles)
+    sin = np.sin(angles)
+    n_pairs = dim // 2
+    out = np.empty_like(x)
+
+    if style == "interleaved":
+        even = x[..., 0::2]
+        odd = x[..., 1::2]
+        cos_b = _broadcast_angles(cos, even)
+        sin_b = _broadcast_angles(sin, odd)
+        out[..., 0::2] = even * cos_b - odd * sin_b
+        out[..., 1::2] = even * sin_b + odd * cos_b
+    else:
+        first = x[..., :n_pairs]
+        second = x[..., n_pairs:]
+        cos_b = _broadcast_angles(cos, first)
+        sin_b = _broadcast_angles(sin, second)
+        out[..., :n_pairs] = first * cos_b - second * sin_b
+        out[..., n_pairs:] = first * sin_b + second * cos_b
+    return out
+
+
+def pair_xy(x: np.ndarray, token: int, pair: int, style: str = "interleaved") -> tuple[float, float]:
+    """Return the 2D coordinates of one RoPE pair at one token."""
+    vec = np.asarray(x)
+    if vec.ndim != 2:
+        raise ValueError("pair_xy expects (seq, dim)")
+    dim = vec.shape[-1]
+    n_pairs = dim // 2
+    if not (0 <= pair < n_pairs):
+        raise IndexError(f"pair {pair} out of range 0..{n_pairs - 1}")
+    if style == "interleaved":
+        return float(vec[token, 2 * pair]), float(vec[token, 2 * pair + 1])
+    if style == "llama":
+        return float(vec[token, pair]), float(vec[token, pair + n_pairs])
+    raise ValueError(f"unknown style {style!r}")
+
+
+def pair_dim_labels(pair: int, dim: int, style: str = "interleaved") -> tuple[str, str]:
+    n_pairs = dim // 2
+    if style == "interleaved":
+        return f"dim {2 * pair}", f"dim {2 * pair + 1}"
+    return f"dim {pair}", f"dim {pair + n_pairs}"
+
+
+def l2_norms(x: np.ndarray) -> np.ndarray:
+    return np.linalg.norm(np.asarray(x, dtype=np.float64), axis=-1)
+
+
+def row_cosine(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    a = np.asarray(a, dtype=np.float64)
+    b = np.asarray(b, dtype=np.float64)
+    an = l2_norms(a)
+    bn = l2_norms(b)
+    dots = np.sum(a * b, axis=-1)
+    return dots / (an * bn + 1e-12)
+
+
+def attention_scores(q: np.ndarray, k: np.ndarray) -> np.ndarray:
+    """Unnormalized ``Q K^T`` for ``(seq, dim)`` slices."""
+    q = np.asarray(q, dtype=np.float64)
+    k = np.asarray(k, dtype=np.float64)
+    return q @ k.T

tests/test_core.py

@@ -0,0 +1,115 @@
+import unittest
+
+import numpy as np
+
+from src.absolute_pe import add_positional_encoding, sinusoidal_pe
+from src.extract import random_qk, select_head
+from src.rope import (
+    apply_rope,
+    attention_scores,
+    l2_norms,
+    pair_frequencies,
+    pair_xy,
+    rotate_pair,
+    row_cosine,
+)
+
+
+class RopeTests(unittest.TestCase):
+    def test_even_dim_required(self):
+        with self.assertRaises(ValueError):
+            apply_rope(np.ones((4, 5)))
+
+    def test_position_zero_is_identity(self):
+        x = np.random.default_rng(0).standard_normal((8, 16))
+        y = apply_rope(x, base=10000.0, style="interleaved")
+        np.testing.assert_allclose(y[0], x[0], atol=1e-12)
+
+    def test_preserves_l2(self):
+        x = np.random.default_rng(1).standard_normal((12, 32))
+        for style in ("interleaved", "llama"):
+            y = apply_rope(x, style=style)
+            np.testing.assert_allclose(l2_norms(x), l2_norms(y), atol=1e-10)
+
+    def test_matches_scratch_even_odd(self):
+        x = np.random.default_rng(2).standard_normal((6, 10))
+        seq, dim = x.shape
+        positions = np.arange(seq)[:, None]
+        pair_indices = np.arange(0, dim, 2)
+        inv_freq = 1 / (10000 ** (pair_indices / dim))
+        angles = positions * inv_freq
+        cos, sin = np.cos(angles), np.sin(angles)
+        expected = np.empty_like(x)
+        expected[:, 0::2] = x[:, 0::2] * cos - x[:, 1::2] * sin
+        expected[:, 1::2] = x[:, 0::2] * sin + x[:, 1::2] * cos
+        np.testing.assert_allclose(apply_rope(x, style="interleaved"), expected)
+
+    def test_relative_angle_depends_on_offset(self):
+        omega = pair_frequencies(8, base=10000.0)[0]
+        x = np.zeros((5, 8))
+        x[:, 0] = 1.0
+        y = apply_rope(x, style="interleaved")
+        xm, ym = pair_xy(y, 3, 0)
+        xn, yn = pair_xy(y, 1, 0)
+        a_m = np.arctan2(ym, xm)
+        a_n = np.arctan2(yn, xn)
+        self.assertAlmostEqual(a_m - a_n, (3 - 1) * omega, places=10)
+
+    def test_rotate_pair_and_batched_heads(self):
+        even, odd = np.array([1.0]), np.array([0.0])
+        re, ro = rotate_pair(even, odd, np.array([np.pi / 2]))
+        np.testing.assert_allclose(re, 0.0, atol=1e-12)
+        np.testing.assert_allclose(ro, 1.0, atol=1e-12)
+        x = np.random.default_rng(3).standard_normal((4, 7, 16))
+        y = apply_rope(x, style="llama")
+        self.assertEqual(y.shape, x.shape)
+        np.testing.assert_allclose(l2_norms(x), l2_norms(y), atol=1e-10)
+
+    def test_cosine_identity_at_zero(self):
+        x = np.random.default_rng(4).standard_normal((5, 12))
+        y = apply_rope(x)
+        np.testing.assert_allclose(row_cosine(x, y)[0], 1.0, atol=1e-10)
+
+    def test_attention_scores_shape(self):
+        q = np.ones((3, 4))
+        k = np.ones((3, 4))
+        s = attention_scores(q, k)
+        self.assertEqual(s.shape, (3, 3))
+
+
+class AbsolutePeTests(unittest.TestCase):
+    def test_matches_loop_formula(self):
+        seq, dim, base = 4, 5, 10000.0
+        pe = sinusoidal_pe(seq, dim, base=base)
+        for k in range(seq):
+            for i in range(dim):
+                if i % 2 == 0:
+                    expected = np.sin(k / base ** (i / dim))
+                else:
+                    expected = np.cos(k / base ** ((i - 1) / dim))
+                self.assertAlmostEqual(pe[k, i], expected, places=12)
+
+    def test_add_changes_norm(self):
+        emb = np.random.default_rng(0).standard_normal((6, 8))
+        pe, combined = add_positional_encoding(emb)
+        self.assertEqual(pe.shape, emb.shape)
+        self.assertFalse(np.allclose(l2_norms(emb), l2_norms(combined)))
+
+
+class ExtractHelpersTests(unittest.TestCase):
+    def test_random_qk(self):
+        data = random_qk(seq_len=8, dim=16, seed=0, base=10000.0)
+        self.assertEqual(data["q_before"].shape, (8, 16))
+        np.testing.assert_allclose(l2_norms(data["q_before"]), l2_norms(data["q_after"]), atol=1e-10)
+        slice_q = select_head(data["q_before"], 99)
+        self.assertEqual(slice_q.shape, (8, 16))
+        from src.plots import bulk_before_after_delta, attention_heatmaps
+
+        fig = bulk_before_after_delta(data["q_before"], data["q_after"])
+        self.assertTrue(len(fig.data) >= 1)
+        s = attention_scores(data["q_before"], data["k_before"])
+        attention_heatmaps(s, s)
+
+
+if __name__ == "__main__":
+    unittest.main()