15 files changed, 957 insertions, 0 deletions
diff --git a/conftest.py b/conftest.py
new file mode 100644
index 0000000..2ca7a99
--- /dev/null
+++ b/conftest.py
@@ -0,0 +1,4 @@
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
diff --git a/design/reviews/gemini-design-review.md b/design/reviews/gemini-design-review.md
new file mode 100644
index 0000000..bc8037c
--- /dev/null
+++ b/design/reviews/gemini-design-review.md
@@ -0,0 +1,63 @@
+### 1. Overall Setup & Feasibility for Next-Prime Grokking
+- **Verdict**: **ISSUE** (High risk of pure memorization without grokking).
+- **Domain size vs. Capacity**: 69 training examples in $[2, 100]$ with ~50k–140k parameters is overparameterized by orders of magnitude. The network can memorize all 69 pairs in a few hundred steps.
+- **Representational Mismatch**: Grokking (Power et al., 2022) typically occurs on algebraic group operations ($x + y \pmod p$) where Fourier representations can compactly solve the task. Next-prime mapped from base-10 digits lacks smooth continuous group symmetries, giving weight decay little structural leverage to discover a compact closed-form algorithm.
+- **Mean-Pooling Blur**: Masked mean-pooling of digit embeddings blurs place value. For example, $10$ and $100$ share token sets $\{1, 0\}$, making exact spatial/positional decoding difficult unless positional embeddings dominate.
+
+---
+
+### 2. Soundness of ACT Setup
+- **Verdict**: **ISSUE** (Flawed state aggregation and penalty schedule).
+- **Non-Standard State Aggregation**: Standard ACT (Graves, 2016; `arXiv:1603.08983`) requires halting weights to sum to 1 ($\sum_{t=1}^N w_t = 1$, using a remainder $R(x)$ at step $N$). Simply taking $\sum p_t h_t$ without normalization allows the state vector norm to scale arbitrarily with step count, causing uncalibrated gradients.
+- **Gradient Shock from Step Warmup**: Hard-switching $\lambda$ from $0$ to $0.01$ at step 1000 introduces a discontinuous loss jump. At batch size 32 on 69 examples (~2 steps/epoch), step 1000 is epoch ~460—by which time training loss is near zero. Turning on $\lambda=0.01$ suddenly causes severe optimization instability.
+- **Better Alternative**: Consider PonderNet (Banino et al., 2021; `arXiv:2107.05407`), which formulates halting probabilistic distributions with a Geometric prior KL penalty, eliminating discrete step-floor collapse. Otherwise, use a smooth linear warmup for $\lambda$ across 5,000 steps.
+
+---
+
+### 3. Optimizer & Weight Decay Tuning
+- **Verdict**: **MIXED** (WD=1.0 is a good baseline; sweep upper bound is too high).
+- **Baseline ($WD = 1.0$)**: **OK**. Strong weight decay is essential to drive network weights out of high-norm memorization regimes into generalizable representations (Power et al., 2022; Nanda et al., 2023).
+- **Sweep Range $\{0.3, 1.0, 3.0, 10.0\}$**: **ISSUE**. For AdamW with $\text{lr}=1\times 10^{-3}$, weight updates decay by $(1 - \eta \cdot \lambda)$ per step. At $\lambda=10.0$, the parameter weight decays by $1\%$ *per step*, which will cause underfitting or complete gradient divergence.
+- **Recommended Sweep**: Shift range to $\{0.01, 0.1, 0.3, 1.0, 3.0\}$ to include lower control baselines.
+
+---
+
+### 4. Prior Art & Theoretical Context
+- **Verdict**: **OK** (Relevant papers identified).
+- **Canonical Grokking Literature**:
+  - Power et al. (2022), *"Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets"*, [`arXiv:2201.02177`](https://arxiv.org/abs/2201.02177).
+  - Nanda et al. (2023), *"Progress measures for grokking via mechanistic interpretability"*, [`arXiv:2301.05217`](https://arxiv.org/abs/2301.05217).
+  - Liu et al. (2022), *"Towards Understanding Grokking: An Empirical Study"*, [`arXiv:2205.10343`](https://arxiv.org/abs/2205.10343).
+  - Varma et al. (2023), *"Explaining Grokking Through Circuit Efficiency"*, [`arXiv:2309.02390`](https://arxiv.org/abs/2309.02390).
+- **Adaptive Computation & Recurrent Looping**:
+  - Graves (2016), *"Adaptive Computation Time for Recurrent Neural Networks"*, [`arXiv:1603.08983`](https://arxiv.org/abs/1603.08983).
+  - Banino et al. (2021), *"PonderNet: Learning to Ponder"*, [`arXiv:2107.05407`](https://arxiv.org/abs/2107.05407).
+  - Giannou et al. (2023), *"Looped Transformers as Programmable Computers"*, [`arXiv:2301.13196`](https://arxiv.org/abs/2301.13196).
+- **Algorithmic Alignment**:
+  - Xu et al. (2020), *"What Can Neural Networks Reason About? Reasoning Algorithmic Alignment"*, [`arXiv:1905.13211`](https://arxiv.org/abs/1905.13211) (ICLR 2020) — shows sample complexity and generalization depend on structural alignment between architecture steps and target algorithm steps (e.g., dynamic programming / trial division steps).
+  - Xu et al. (2021), *"How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks"*, [`arXiv:2009.11848`](https://arxiv.org/abs/2009.11848).
+- **Primality Domain Note**: Past grokking work focuses on *prime modulo arithmetic* ($x+y \pmod p$), not predicting *next-prime* from base-10 representations.
+
+---
+
+### 5. GRU Decoder Architectural Confound
+- **Verdict**: **ISSUE** (Breaks strict weight-tying comparison).
+- **Confound**: Introducing an un-tied 1-layer GRU decoder on Arm A creates an architectural asymmetry against Arm B (GPT-style causal transformer). The GRU decoder contains enough un-tied parameters to perform sequential pattern lookup on output tokens, masking whether the tied recurrent cell actually solved the problem.
+- **Cleaner Alternatives**:
+  1. **Direct Tied Linear Readout**: Project the final state $h_{\text{final}}$ directly to vocabulary logits via a single linear layer (or shared embedding matrix).
+  2. **Recurrent Cell Decoding**: Feed previous output tokens back into the *same tied cell* during generation rather than using a separate GRU block.
+
+---
+
+### 6. Interpretation Logic & Metrics
+- **Verdict**: **MIXED** (Valid in-domain criteria, but needs extrapolation check).
+- **In-Domain Sharp Transition**: Train $\text{EM} \ge 0.95 \to$ Val $\text{EM}$ jumping $0.2 \to 0.9$ within 5 evals correctly identifies grokking dynamics on $[2, 100]$.
+- **Out-of-Domain Generalization ($[101, 200]$)**: High val EM on $[2, 100]$ alone does *not* prove a primality algorithm (sieve/trial division) was grokked. If validation accuracy on $[101, 200]$ remains near 0%, the network merely grokked a bounded interpolation table for $n \le 100$.
+
+---
+
+### Top 3 Things to Fix Before Launching Runs
+
+1. **Remove the GRU Decoder Confound**: Replace the un-tied GRU decoder in Arm A with a single shared linear projection layer to strictly isolate the effect of weight-tied recurrence.
+2. **Fix the ACT Formulation & Warmup Schedule**: Normalize state weighting ($\sum w_t = 1$) following Graves (2016) or switch to PonderNet (`arXiv:2107.05407`). Replace the abrupt step-1000 penalty turn-on with a smooth linear schedule for $\lambda$ across 5,000 steps.
+3. **Fix Input Representation & Alignment**: Replace mean-pooling with explicit sequence concatenation/causal prefix encoding to preserve place-value structure. Lower the AdamW weight decay sweep upper bound from $10.0$ to $3.0$.
diff --git a/scripts/plot.py b/scripts/plot.py
new file mode 100644
index 0000000..14536e5
--- /dev/null
+++ b/scripts/plot.py
@@ -0,0 +1,49 @@
+"""Plot train/val curves from a metrics.csv. Usage: python -m scripts.plot <model> <seed>"""
+import csv
+import sys
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+
+def main() -> None:
+    model, seed = sys.argv[1], sys.argv[2]
+    path = f"runs/{model}/seed{seed}/metrics.csv"
+    with open(path) as fh:
+        rows = list(csv.DictReader(fh))
+    steps = [int(r["step"]) for r in rows]
+    train_loss = [float(r["train_loss"]) for r in rows]
+    train_em = [float(r["train_em"]) for r in rows]
+    val_em = [float(r["val_em"]) for r in rows]
+    train_tok = [float(r["train_token_acc"]) for r in rows]
+    val_tok = [float(r["val_token_acc"]) for r in rows]
+    halt = [float(r["mean_halt_steps"]) for r in rows]
+
+    fig, axes = plt.subplots(2, 2, figsize=(12, 8))
+    axes[0, 0].plot(steps, train_loss, label="train loss", color="tab:blue")
+    axes[0, 0].set_title("Train loss (token CE + halt penalty)")
+    axes[0, 0].set_xlabel("step")
+    axes[0, 1].plot(steps, train_em, label="train EM", color="tab:orange")
+    axes[0, 1].plot(steps, val_em, label="val EM", color="tab:green")
+    axes[0, 1].axhline(0.9, ls="--", c="gray", lw=0.7)
+    axes[0, 1].set_title(f"Exact-match (model={model}, seed={seed})")
+    axes[0, 1].set_ylim(-0.05, 1.05)
+    axes[0, 1].legend()
+    axes[1, 0].plot(steps, train_tok, label="train tok acc", color="tab:red")
+    axes[1, 0].plot(steps, val_tok, label="val tok acc", color="tab:purple")
+    axes[1, 0].set_title("Token accuracy")
+    axes[1, 0].set_ylim(-0.05, 1.05)
+    axes[1, 0].legend()
+    axes[1, 1].plot(steps, halt, label="mean halt steps", color="tab:brown")
+    axes[1, 1].set_title("RNN halting (mean steps used)")
+    axes[1, 1].set_xlabel("step")
+    fig.suptitle(f"prime-grokking — {model} seed {seed}")
+    plt.tight_layout()
+    out = f"runs/{model}/seed{seed}/curves.png"
+    plt.savefig(out, dpi=110)
+    print(f"saved {out}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/run_experiment.sh b/scripts/run_experiment.sh
new file mode 100755
index 0000000..feaef87
--- /dev/null
+++ b/scripts/run_experiment.sh
@@ -0,0 +1,4 @@
+#!/usr/bin/env bash
+set -euo pipefail
+cd "$(dirname "$0")/.."
+exec .venv/bin/python -m src.train "$@"
diff --git a/src/__init__.py b/src/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/src/__init__.py
diff --git a/src/config.py b/src/config.py
new file mode 100644
index 0000000..74cb4d3
--- /dev/null
+++ b/src/config.py
@@ -0,0 +1,101 @@
+"""Experiment configuration. Defaults = Experiment 1 (design/preregistration.md)."""
+import argparse
+import json
+from dataclasses import dataclass, fields
+
+
+@dataclass
+class Config:
+    # data
+    vocab_mode: str = "digits"       # "digits" | "integers"
+    range_start: int = 2
+    range_end: int = 100             # inclusive
+    holdout_frac: float = 0.30
+    seed: int = 0
+    # model
+    model: str = "rnn"               # "rnn" | "transformer"
+    d_model: int = 128
+    max_steps: int = 20              # K tied iterations (RNN)
+    halting: bool = True             # False -> fixed-K ablation
+    halt_eps: float = 0.05           # ACT cumulative threshold (documented; K small so no early break)
+    halt_penalty: float = 0.01       # lambda on mean steps
+    halt_warmup_steps: int = 1000    # penalty off before this (anti-collapse)
+    min_steps: int = 2               # ACT: halt prob forced to 0 before this step
+    n_layers: int = 2
+    n_heads: int = 4
+    # training
+    lr: float = 1e-3
+    weight_decay: float = 1.0
+    max_train_steps: int = 200_000
+    eval_every: int = 200
+    early_stop_em: float = 1.0
+    early_stop_patience: int = 5
+    batch_size: int = 32
+    max_out_len: int = 6
+    log_n_examples: int = 10
+    out_dir: str = "runs"
+
+    @property
+    def vocab(self) -> int:
+        """Token count. Integers mode: 0..range_end+1 (next prime can exceed range_end)."""
+        if self.vocab_mode == "integers":
+            return self.range_end + 2
+        return 11                      # digits 0-9 + EOS
+
+    @property
+    def eos_id(self) -> int:
+        return self.vocab - 1
+
+    @property
+    def pad_id(self) -> int:
+        return self.vocab              # one extra embedding row reserved for pad
+
+    def to_json(self) -> dict:
+        d = {f.name: getattr(self, f.name) for f in fields(self)}
+        d["vocab"] = self.vocab
+        d["eos_id"] = self.eos_id
+        d["pad_id"] = self.pad_id
+        return d
+
+    @classmethod
+    def from_json(cls, d: dict) -> "Config":
+        cfg = cls()
+        for f in fields(cls):
+            if f.name in d:
+                setattr(cfg, f.name, d[f.name])
+        return cfg
+
+    def save(self, path: str) -> None:
+        with open(path, "w") as fh:
+            json.dump(self.to_json(), fh, indent=2)
+
+    @classmethod
+    def load(cls, path: str) -> "Config":
+        with open(path) as fh:
+            return cls.from_json(json.load(fh))
+
+
+def _bool_arg(s: str) -> bool:
+    return s.lower() in ("1", "true", "yes", "on")
+
+
+def parse_args(argv=None) -> Config:
+    cfg = Config()
+    p = argparse.ArgumentParser(description="prime-grokking train")
+    p.add_argument("model_pos", nargs="?", default=None, help="model: rnn | transformer")
+    p.add_argument("seed_pos", nargs="?", default=None, help="seed (int)")
+    for f in fields(Config):
+        if f.type is bool:
+            p.add_argument(f"--{f.name}", default=None, type=_bool_arg)
+        elif f.type in (int, float, str):
+            p.add_argument(f"--{f.name}", default=None, type=f.type)
+    a = p.parse_args(argv)
+    if a.model_pos is not None:
+        cfg.model = a.model_pos
+    if a.seed_pos is not None:
+        cfg.seed = int(a.seed_pos)
+    for f in fields(Config):
+        v = getattr(a, f.name, None)
+        if v is not None:
+            setattr(cfg, f.name, v)
+    return cfg
diff --git a/src/data.py b/src/data.py
new file mode 100644
index 0000000..35682d3
--- /dev/null
+++ b/src/data.py
@@ -0,0 +1,79 @@
+"""Prime dataset: n -> next prime, digit-tokenized; splits and batching."""
+import random
+
+import torch
+
+from src.config import Config
+
+
+def sieve_primes(limit: int) -> list[int]:
+    """All primes <= limit (inclusive)."""
+    if limit < 2:
+        return []
+    is_prime = [True] * (limit + 1)
+    is_prime[0] = is_prime[1] = False
+    for p in range(2, int(limit ** 0.5) + 1):
+        if is_prime[p]:
+            for m in range(p * p, limit + 1, p):
+                is_prime[m] = False
+    return [i for i in range(2, limit + 1) if is_prime[i]]
+
+
+def next_prime(n: int, primes: list[int]) -> int:
+    for p in primes:
+        if p > n:
+            return p
+    raise ValueError(f"no prime > {n} in supplied list")
+
+
+def encode_int(n: int, cfg: Config) -> list[int]:
+    if cfg.vocab_mode == "integers":
+        return [n]
+    return [int(d) for d in str(n)]
+
+
+def decode_tokens(ts, cfg: Config) -> int:
+    """Decode a token sequence, stopping at EOS. -1 if nothing decodable."""
+    if cfg.vocab_mode == "integers":
+        return int(ts[0]) if len(ts) else -1
+    digits = []
+    for t in ts:
+        if t == cfg.eos_id:
+            break
+        digits.append(int(t))
+    return int("".join(map(str, digits))) if digits else -1
+
+
+def get_splits(cfg: Config) -> tuple[list[int], list[int]]:
+    """(train, val) input lists, seeded shuffle, no overlap."""
+    rng = random.Random(cfg.seed)
+    inputs = list(range(cfg.range_start, cfg.range_end + 1))
+    rng.shuffle(inputs)
+    n_val = max(1, round(len(inputs) * cfg.holdout_frac))
+    return sorted(inputs[n_val:]), sorted(inputs[:n_val])
+
+
+def build_examples(inputs: list[int], cfg: Config) -> list[tuple[list[int], list[int]]]:
+    """[(input_tokens, target_tokens+EOS), ...]"""
+    primes = sieve_primes(cfg.range_end + 100)
+    out = []
+    for n in inputs:
+        p = next_prime(n, primes)
+        out.append((encode_int(n, cfg), encode_int(p, cfg) + [cfg.eos_id]))
+    return out
+
+
+def make_batch(examples, cfg: Config) -> dict[str, torch.Tensor]:
+    xs, ys = zip(*examples)
+    T_in = max(len(x) for x in xs)
+    T_out = max(len(y) for y in ys)
+    B = len(examples)
+    x = torch.full((B, T_in), cfg.pad_id, dtype=torch.long)
+    y = torch.full((B, T_out), cfg.pad_id, dtype=torch.long)
+    for i, (xi, yi) in enumerate(examples):
+        x[i, : len(xi)] = torch.tensor(xi, dtype=torch.long)
+        y[i, : len(yi)] = torch.tensor(yi, dtype=torch.long)
+    # teacher-forced decoder input: BOS(=EOS reuse) then shifted y
+    y_in = torch.cat([torch.full((B, 1), cfg.eos_id, dtype=torch.long), y[:, :-1]], dim=1)
+    y_mask = y != cfg.pad_id
+    return {"x": x, "y": y, "y_in": y_in, "y_mask": y_mask}
diff --git a/src/eval.py b/src/eval.py
new file mode 100644
index 0000000..9d55437
--- /dev/null
+++ b/src/eval.py
@@ -0,0 +1,151 @@
+"""Post-run analysis: final metrics, [101,200] probe with sieve diagnostic, grokking signature.
+
+Interpretation codes are locked in design/preregistration.md — this module only MEASURES
+and classifies against those definitions (O1-O4, H1-H4, P1-P4).
+"""
+import argparse
+import csv
+import json
+import os
+
+import numpy as np
+import torch
+
+from src.config import Config
+from src.data import build_examples, decode_tokens, encode_int, get_splits, next_prime, sieve_primes
+from src.model_api import build_model, greedy_decode
+from src.train import evaluate
+
+FLAGGED_COMPOSITES = {121, 143, 169, 187}   # need divisors 11, 13 — beyond the {2,3,5,7} sieve
+
+
+def probe_report(model, cfg: Config, lo: int = 101, hi: int = 200) -> dict:
+    primes = sieve_primes(hi + 200)
+    correct = 0
+    errors = []
+    easy_misses = 0
+    for n in range(lo, hi + 1):
+        x = torch.tensor(encode_int(n, cfg), dtype=torch.long).unsqueeze(0)
+        gen = greedy_decode(model, x, cfg)[0].tolist()
+        pred = decode_tokens(gen, cfg)
+        target = next_prime(n, primes)
+        if pred == target:
+            correct += 1
+        else:
+            errors.append({"n": n, "target": target, "pred": pred})
+            if n % 2 == 0 or n % 5 == 0:
+                easy_misses += 1
+    total = hi - lo + 1
+    acc = correct / total
+    flagged = [e for e in errors if e["n"] in FLAGGED_COMPOSITES]
+    # P-code classification (see preregistration.md)
+    if acc >= 0.85:
+        code = "P3"
+    elif easy_misses >= 5:
+        code = "P4"
+    elif errors and all(e["n"] in FLAGGED_COMPOSITES for e in errors) and len(errors) <= 6:
+        code = "P1"
+    else:
+        code = "P2"
+    return {
+        "code": code, "acc": acc, "correct": correct, "total": total,
+        "errors": errors, "flagged_errors": flagged, "easy_misses": easy_misses,
+    }
+
+
+def grokking_signature(metrics_path: str) -> dict:
+    """Classify the training curve against preregistered codes O1-O4."""
+    with open(metrics_path) as fh:
+        rows = list(csv.DictReader(fh))
+    if not rows:
+        return {"code": "O4", "note": "no eval rows"}
+    train = [float(r["train_em"]) for r in rows]
+    val = [float(r["val_em"]) for r in rows]
+    n = len(rows)
+    saturated = any(all(t >= 0.95 for t in train[i:i + 10]) for i in range(n - 9)) if n >= 10 else False
+    hi = next((i for i, v in enumerate(val) if v >= 0.9), None)
+    trans = None
+    if hi is not None:
+        lo_cands = [i for i in range(hi) if val[i] <= 0.2]
+        if lo_cands:
+            trans = hi - max(lo_cands)
+    if saturated and hi is not None and trans is not None and trans <= 5:
+        code = "O1"
+    elif max(train) >= 0.95 and hi is not None and (trans is None or trans > 5):
+        code = "O3"
+    elif max(train) >= 0.95 and hi is None:
+        code = "O2"
+    else:
+        code = "O4"
+    return {
+        "code": code, "train_saturated": saturated, "val_hi_eval_idx": hi,
+        "transition_width_evals": trans, "n_evals": n,
+        "final_train_em": train[-1], "final_val_em": val[-1],
+    }
+
+
+@torch.no_grad()
+def halting_report(model, cfg: Config) -> dict:
+    """RNN halting structure: mean steps + correlation with gap-to-next-prime (H1-H4)."""
+    primes = sieve_primes(300)
+    gaps, steps = [], []
+    for n in range(cfg.range_start, cfg.range_end + 1):
+        x = torch.tensor(encode_int(n, cfg), dtype=torch.long).unsqueeze(0)
+        h = model._initial_state(x)
+        _, s = model._run_cell(h)
+        gaps.append(next_prime(n, primes) - n)
+        steps.append(float(s.mean()))
+    mean = float(np.mean(steps))
+    rho = float(np.corrcoef(gaps, steps)[0, 1]) if len(set(gaps)) > 1 else 0.0
+    if mean <= cfg.min_steps + 0.5:
+        code = "H1"
+    elif mean >= cfg.max_steps - 0.5:
+        code = "H2"
+    elif abs(rho) >= 0.3:
+        code = "H3"
+    else:
+        code = "H4"
+    return {"code": code, "mean_steps": mean, "corr_gap": rho, "min_steps": cfg.min_steps, "max_steps": cfg.max_steps}
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description="prime-grokking eval")
+    ap.add_argument("model")
+    ap.add_argument("seed")
+    ap.add_argument("--ckpt", default="best.pt")
+    a = ap.parse_args()
+    out_dir = os.path.join("runs", a.model, f"seed{a.seed}")
+    cfg = Config.load(os.path.join(out_dir, "config.json"))
+    model = build_model(cfg)
+    model.load_state_dict(torch.load(os.path.join(out_dir, a.ckpt), map_location="cpu"))
+    model.eval()
+
+    _, val_in = get_splits(cfg)
+    val_ex = build_examples(val_in, cfg)
+    tok, em, per = evaluate(model, val_ex, cfg)
+
+    probe = probe_report(model, cfg)
+    sig = grokking_signature(os.path.join(out_dir, "metrics.csv"))
+    hlt = halting_report(model, cfg) if cfg.model == "rnn" else None
+
+    results = {
+        "model": cfg.model, "seed": cfg.seed, "ckpt": a.ckpt,
+        "params": model.param_count(),
+        "val_token_acc": tok, "val_exact_match": em,
+        "per_example": [{"n": "".join(map(str, x)), "target": t, "pred": p, "ok": ok}
+                        for x, t, p, ok in per],
+        "probe": probe,
+        "signature": sig,
+        "halting": hlt,
+    }
+    with open(os.path.join(out_dir, "results.json"), "w") as fh:
+        json.dump(results, fh, indent=2)
+    print(json.dumps({"val_token_acc": round(tok, 4), "val_em": round(em, 4),
+                      "probe": probe["code"], "probe_acc": round(probe["acc"], 3),
+                      "flagged_errors": [e["n"] for e in probe["flagged_errors"]],
+                      "signature": sig["code"], "halting": hlt["code"] if hlt else None,
+                      "results": os.path.join(out_dir, "results.json")}, indent=2))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/model_api.py b/src/model_api.py
new file mode 100644
index 0000000..66c2ace
--- /dev/null
+++ b/src/model_api.py
@@ -0,0 +1,38 @@
+"""Model interface contract + build dispatch + shared greedy decode."""
+import torch
+import torch.nn as nn
+
+from src.config import Config
+
+
+class PrimeModel(nn.Module):
+    """Contract: forward(x, y_in) -> {"logits": (B,T_out,vocab), "halt_steps": (B,) or None}."""
+
+    def forward(self, x, y_in):
+        raise NotImplementedError
+
+    def param_count(self) -> int:
+        return sum(p.numel() for p in self.parameters())
+
+
+def build_model(cfg: Config) -> PrimeModel:
+    if cfg.model == "rnn":
+        from src.models.rnn import TiedRNN
+        return TiedRNN(cfg)
+    if cfg.model == "transformer":
+        from src.models.transformer import TransformerBaseline
+        return TransformerBaseline(cfg)
+    raise ValueError(f"unknown model: {cfg.model}")
+
+
+@torch.no_grad()
+def greedy_decode(model: PrimeModel, x: torch.Tensor, cfg: Config, max_len: int | None = None) -> torch.Tensor:
+    """Autoregressive greedy decode of output digits. Returns (B, max_len) tokens (BOS stripped)."""
+    max_len = max_len or cfg.max_out_len
+    B = x.shape[0]
+    y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=x.device)
+    for _ in range(max_len):
+        out = model(x, y_in)
+        nxt = out["logits"][:, -1].argmax(-1)
+        y_in = torch.cat([y_in, nxt[:, None]], dim=1)
+    return y_in[:, 1:]
diff --git a/src/models/__init__.py b/src/models/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/src/models/__init__.py
diff --git a/src/models/rnn.py b/src/models/rnn.py
new file mode 100644
index 0000000..9947c66
--- /dev/null
+++ b/src/models/rnn.py
@@ -0,0 +1,98 @@
+"""Weight-tied RNN: one 2-layer cell applied K times, ACT learned halting, GRU digit decoder.
+
+Spec pseudocode (design/experiment-spec.md):
+    state = embed(input_number)
+    for step in range(max_steps):
+        state = step_module(state)      # same weights every iteration
+        if halt_condition(state): break
+    output = project(state)
+
+Initial state: masked mean-pool of (digit embedding + sinusoidal positional encoding)
+passed through a small MLP, so digit ORDER reaches the tied cell.
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from src.config import Config
+from src.model_api import PrimeModel
+
+
+class TiedRNN(PrimeModel):
+    def __init__(self, cfg: Config):
+        super().__init__()
+        self.cfg = cfg
+        d = cfg.d_model
+        self.embed = nn.Embedding(cfg.vocab + 1, d)      # +1 row = pad
+        self.in_proj = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, d))
+        self.ln1 = nn.LayerNorm(d)
+        self.cell_ln = nn.LayerNorm(d)
+        self.cell_w1 = nn.Linear(d, d)
+        self.cell_w2 = nn.Linear(d, d)
+        self.halt_head = nn.Linear(d, 1)
+        self.decoder = nn.GRUCell(d, d)
+        self.out_head = nn.Linear(d, cfg.vocab)
+
+    @staticmethod
+    def _sinusoidal(T: int, d: int) -> torch.Tensor:
+        pe = torch.zeros(T, d)
+        pos = torch.arange(T).float().unsqueeze(1)
+        i = torch.arange(d).float().unsqueeze(0)
+        pe[:, 0::2] = torch.sin(pos / 10000 ** (2 * i[:, 0::2] / d))
+        pe[:, 1::2] = torch.cos(pos / 10000 ** (2 * i[:, 1::2] / d))
+        return pe
+
+    def _initial_state(self, x: torch.Tensor) -> torch.Tensor:
+        B, T = x.shape
+        mask = (x != self.cfg.pad_id).float().unsqueeze(-1)          # (B,T,1)
+        pos = self._sinusoidal(T, self.cfg.d_model).to(x.device)     # (T,d)
+        e = self.embed(x) + pos.unsqueeze(0)                          # (B,T,d)
+        h = (e * mask).sum(1) / mask.sum(1).clamp(min=1)             # (B,d)
+        return self.in_proj(h)
+
+    def _run_cell(self, h0: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Tied cell x K steps. Returns (final_state (B,d), mean_steps (B,))."""
+        cfg = self.cfg
+        B = h0.shape[0]
+        device = h0.device
+        if not cfg.halting:
+            h = h0
+            for _ in range(cfg.max_steps):
+                h = h + self.cell_w2(F.gelu(self.cell_w1(self.cell_ln(h))))
+            steps = h0.new_full((B,), float(cfg.max_steps))
+            return h, steps
+        # ACT: run all K steps, accumulate weighted average (K=20 -> no early break needed)
+        h_list, p_list = [], []
+        h = h0
+        for t in range(cfg.max_steps):
+            h = h + self.cell_w2(F.gelu(self.cell_w1(self.cell_ln(h))))
+            p = torch.sigmoid(self.halt_head(self.ln1(h))).squeeze(-1)   # (B,)
+            if t < cfg.min_steps:
+                p = p * 0.0
+            h_list.append(h)
+            p_list.append(p)
+        final = torch.zeros_like(h0)
+        steps = torch.zeros(B, device=device)
+        remaining = torch.ones(B, device=device)
+        for t in range(cfg.max_steps):
+            p = p_list[t]
+            w = remaining * p
+            final = final + w.unsqueeze(-1) * h_list[t]
+            steps = steps + (t + 1) * w
+            remaining = remaining * (1 - p)
+        final = final + remaining.unsqueeze(-1) * h_list[-1]
+        steps = steps + remaining * cfg.max_steps
+        return final, steps
+
+    def forward(self, x: torch.Tensor, y_in: torch.Tensor) -> dict:
+        cfg = self.cfg
+        h = self._initial_state(x)
+        h, steps = self._run_cell(h)
+        # autoregressive digit decoder, teacher-forced during training
+        e = self.embed(y_in)                       # (B,T_out,d)
+        outs = []
+        for t in range(y_in.shape[1]):
+            h = self.decoder(e[:, t], h)
+            outs.append(self.out_head(h))
+        logits = torch.stack(outs, dim=1)          # (B,T_out,vocab)
+        return {"logits": logits, "halt_steps": steps}
diff --git a/src/models/transformer.py b/src/models/transformer.py
new file mode 100644
index 0000000..32e158b
--- /dev/null
+++ b/src/models/transformer.py
@@ -0,0 +1,61 @@
+"""Transformer baseline: GPT-style causal decoder over [input digits | output digits].
+
+Fixed d_model (=128) matching the RNN arm. Parameter counts are logged per run but
+NOT gated to parity (design/preregistration.md: weight sharing is the studied variable).
+"""
+import torch
+import torch.nn as nn
+
+from src.config import Config
+from src.model_api import PrimeModel
+
+
+class CausalBlock(nn.Module):
+    def __init__(self, d: int, heads: int):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(d)
+        self.ln2 = nn.LayerNorm(d)
+        self.attn = nn.MultiheadAttention(d, heads, batch_first=True)
+        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
+
+    def forward(self, x: torch.Tensor, attn_mask: torch.Tensor, key_pad: torch.Tensor) -> torch.Tensor:
+        a, _ = self.attn(self.ln1(x), self.ln1(x), self.ln1(x),
+                         attn_mask=attn_mask, key_padding_mask=key_pad, need_weights=False)
+        x = x + a
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+
+class TransformerBaseline(PrimeModel):
+    def __init__(self, cfg: Config):
+        super().__init__()
+        self.cfg = cfg
+        d = cfg.d_model
+        self.embed = nn.Embedding(cfg.vocab + 1, d)     # +1 row = pad
+        self.pos = nn.Embedding(64, d)                  # 3 in + 6 out max, generous
+        self.blocks = nn.ModuleList([CausalBlock(d, cfg.n_heads) for _ in range(cfg.n_layers)])
+        self.ln_f = nn.LayerNorm(d)
+        self.head = nn.Linear(d, cfg.vocab)
+
+    def forward(self, x: torch.Tensor, y_in: torch.Tensor) -> dict:
+        cfg = self.cfg
+        B, T_in = x.shape
+        T_out = y_in.shape[1]
+        seq = torch.cat([x, y_in], dim=1)               # (B, T_in+T_out)
+        T = seq.shape[1]
+        device = seq.device
+        e = self.embed(seq) + self.pos(torch.arange(T, device=device)).unsqueeze(0)
+        # causal mask: input positions (j < T_in) fully visible; output positions causal
+        # (True = blocked, per torch.nn.MultiheadAttention bool convention)
+        blocked = torch.zeros(T, T, dtype=torch.bool, device=device)
+        for i in range(T):
+            for j in range(T):
+                if j >= T_in and j > i:
+                    blocked[i, j] = True
+        key_pad = seq == cfg.pad_id                       # (B,T) True = ignore
+        h = e
+        for blk in self.blocks:
+            h = blk(h, blocked, key_pad)
+        # logits at position p predict token p+1 -> positions [T_in-1, T_in+T_out-2] predict y
+        logits = self.head(self.ln_f(h))[:, T_in - 1: T_in - 1 + T_out]
+        return {"logits": logits, "halt_steps": None}
diff --git a/src/train.py b/src/train.py
new file mode 100644
index 0000000..f0c63e9
--- /dev/null
+++ b/src/train.py
@@ -0,0 +1,154 @@
+"""Training loop. Usage: python -m src.train [model] [seed] [--flag ...]"""
+import csv
+import math
+import os
+import random
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from src.config import Config, parse_args
+from src.data import build_examples, decode_tokens, get_splits, make_batch
+from src.model_api import build_model, greedy_decode
+
+
+def set_seed(s: int) -> None:
+    random.seed(s)
+    np.random.seed(s)
+    torch.manual_seed(s)
+
+
+@torch.no_grad()
+def evaluate(model, examples, cfg: Config):
+    """Token accuracy (teacher-forced) + exact-match accuracy (greedy) + per-example detail."""
+    model.eval()
+    token_correct = 0
+    token_total = 0
+    em_correct = 0
+    per_example = []
+    for x, target in examples:
+        batch = make_batch([(x, target)], cfg)
+        out = model(batch["x"], batch["y_in"])
+        logits = out["logits"][0]
+        y = batch["y"][0]
+        mask = batch["y_mask"][0]
+        pred_tok = logits.argmax(-1)
+        token_correct += int((pred_tok[mask] == y[mask]).sum())
+        token_total += int(mask.sum())
+        xi = torch.tensor(x, dtype=torch.long).unsqueeze(0)
+        gen = greedy_decode(model, xi, cfg)[0].tolist()
+        pred = decode_tokens(gen, cfg)
+        target_n = decode_tokens(target, cfg)
+        ok = pred == target_n
+        em_correct += int(ok)
+        per_example.append((tuple(x), target_n, pred, ok))
+    model.train()
+    return token_correct / max(1, token_total), em_correct / len(examples), per_example
+
+
+def log_example_rows(log_examples, val_per_example):
+    """Format the fixed 10 logged val inputs as '42:ok' strings, in fixed order."""
+    by_x = {x: (target, pred, ok) for x, target, pred, ok in val_per_example}
+    rows = []
+    for x, _target in log_examples:
+        t, p, ok = by_x[tuple(x)]
+        label = "".join(map(str, x))  # "42" in both vocab modes
+        rows.append(f"{label}:{'ok' if ok else f'{p}~{t}'}")
+    return ";".join(rows)
+
+
+def main() -> None:
+    cfg = parse_args()
+    set_seed(cfg.seed)
+    out_dir = os.path.join(cfg.out_dir, cfg.model, f"seed{cfg.seed}")
+    os.makedirs(out_dir, exist_ok=True)
+    cfg.save(os.path.join(out_dir, "config.json"))
+
+    train_in, val_in = get_splits(cfg)
+    train_ex = build_examples(train_in, cfg)
+    val_ex = build_examples(val_in, cfg)
+    model = build_model(cfg)
+    print(f"model={cfg.model} params={model.param_count()} train={len(train_ex)} val={len(val_ex)} "
+          f"vocab={cfg.vocab} eos={cfg.eos_id} pad={cfg.pad_id}")
+
+    opt = torch.optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)
+    ce = nn.CrossEntropyLoss(reduction="none")
+
+    log_examples = sorted(val_ex, key=lambda x: (len(str(x)), x))[: cfg.log_n_examples]
+
+    csv_path = os.path.join(out_dir, "metrics.csv")
+    fieldnames = ["step", "train_loss", "train_token_acc", "train_em", "val_token_acc",
+                  "val_em", "mean_halt_steps", "log_examples", "param_count"]
+    best_val_em = -1.0
+    patience_left = cfg.early_stop_patience
+    best_path = os.path.join(out_dir, "best.pt")
+
+    n_batches = math.ceil(len(train_ex) / cfg.batch_size)
+    rng = random.Random(cfg.seed)
+    step = 0
+    done = False
+    first_row = True
+
+    def _eval_pass(cur_loss, halt_steps):
+        nonlocal best_val_em, patience_left, done, first_row
+        train_tok, train_em, _ = evaluate(model, train_ex, cfg)
+        val_tok, val_em, val_per = evaluate(model, val_ex, cfg)
+        mean_halt = float(halt_steps.mean()) if halt_steps is not None else float("nan")
+        row = {
+            "step": step, "train_loss": float(cur_loss), "train_token_acc": train_tok,
+            "train_em": train_em, "val_token_acc": val_tok, "val_em": val_em,
+            "mean_halt_steps": mean_halt,
+            "log_examples": log_example_rows(log_examples, val_per),
+            "param_count": model.param_count(),
+        }
+        with open(csv_path, "a", newline="") as fh:
+            w = csv.DictWriter(fh, fieldnames=fieldnames)
+            if first_row:
+                w.writeheader()
+                first_row = False
+            w.writerow(row)
+        if val_em > best_val_em:
+            best_val_em = val_em
+            torch.save(model.state_dict(), best_path)
+        if val_em >= cfg.early_stop_em:
+            patience_left -= 1
+        else:
+            patience_left = cfg.early_stop_patience
+        if patience_left <= 0:
+            done = True
+        print(f"step {step} loss {float(cur_loss):.4f} train_em {train_em:.3f} "
+              f"val_em {val_em:.3f} val_tok {val_tok:.3f} halt {mean_halt:.2f}")
+
+    while step < cfg.max_train_steps and not done:
+        rng.shuffle(train_ex)
+        for bi in range(n_batches):
+            sl = train_ex[bi * cfg.batch_size: (bi + 1) * cfg.batch_size]
+            if not sl:
+                continue
+            batch = make_batch(sl, cfg)
+            out = model(batch["x"], batch["y_in"])
+            logits = out["logits"]
+            y_safe = batch["y"].clamp(max=cfg.vocab - 1)   # CE index guard for pad positions
+            loss_tokens = ce(logits.reshape(-1, cfg.vocab), y_safe.reshape(-1)).reshape(
+                logits.shape[0], -1) * batch["y_mask"].float()
+            loss_tokens = loss_tokens.sum() / batch["y_mask"].sum().clamp(min=1)
+            lam = cfg.halt_penalty if (cfg.halting and step >= cfg.halt_warmup_steps) else 0.0
+            halt = out["halt_steps"]
+            penalty = halt.float().mean() * lam if halt is not None and lam > 0 else 0.0
+            loss = loss_tokens + penalty
+            opt.zero_grad()
+            loss.backward()
+            opt.step()
+            step += 1
+            if step % cfg.eval_every == 0 or step >= cfg.max_train_steps:
+                _eval_pass(loss.detach(), halt.detach() if halt is not None else None)
+                if done:
+                    break
+
+    torch.save(model.state_dict(), os.path.join(out_dir, "last.pt"))
+    print(f"DONE steps={step} best_val_em={best_val_em:.4f}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/test_data.py b/tests/test_data.py
new file mode 100644
index 0000000..fe84e94
--- /dev/null
+++ b/tests/test_data.py
@@ -0,0 +1,78 @@
+import torch
+
+from src.config import Config
+from src.data import (build_examples, decode_tokens, encode_int, get_splits,
+                      make_batch, next_prime, sieve_primes)
+
+
+def test_sieve_primes_known():
+    assert sieve_primes(30) == [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
+    assert sieve_primes(1) == []
+    assert sieve_primes(2) == [2]
+
+
+def test_next_prime():
+    p = sieve_primes(200)
+    assert next_prime(42, p) == 43
+    assert next_prime(100, p) == 101
+    assert next_prime(2, p) == 3
+    assert next_prime(97, p) == 101
+
+
+def test_encode_decode_roundtrip_digits():
+    cfg = Config(vocab_mode="digits")
+    for n in [2, 7, 10, 42, 99, 100]:
+        assert decode_tokens(encode_int(n, cfg) + [cfg.eos_id], cfg) == n
+    assert encode_int(42, cfg) == [4, 2]
+
+
+def test_encode_decode_roundtrip_integers():
+    cfg = Config(vocab_mode="integers", range_end=100)
+    for n in [2, 42, 100]:
+        assert decode_tokens(encode_int(n, cfg) + [cfg.eos_id], cfg) == n
+    assert cfg.vocab == 102  # 0..101 + EOS
+
+
+def test_splits_sizes_and_no_overlap():
+    cfg = Config()
+    tr, va = get_splits(cfg)
+    assert len(tr) == 69 and len(va) == 30
+    assert not set(tr) & set(va)
+    assert set(tr) | set(va) == set(range(2, 101))
+
+
+def test_splits_seed_stable():
+    a1, b1 = get_splits(Config(seed=3))
+    a2, b2 = get_splits(Config(seed=3))
+    assert a1 == a2 and b1 == b2
+    a3, _ = get_splits(Config(seed=4))
+    assert a1 != a3
+
+
+def test_build_examples_targets():
+    cfg = Config()
+    ex = build_examples([2, 42, 99, 100], cfg)
+    assert [decode_tokens(y, cfg) for _, y in ex] == [3, 43, 101, 101]
+    assert ex[0][1][-1] == cfg.eos_id  # EOS-terminated
+
+
+def test_make_batch_shapes_and_padding():
+    cfg = Config()
+    ex = build_examples([2, 7, 42, 99], cfg)
+    b = make_batch(ex, cfg)
+    assert b["x"].shape == (4, 2)          # max input len 2 (42, 99)
+    assert b["y"].shape == (4, 4)          # max output len 4 ("101"+EOS)
+    assert b["y_in"].shape == (4, 4)
+    assert b["y_mask"].shape == (4, 4)
+    # pad positions masked out
+    assert not b["y_mask"][0, 3]           # "3"+EOS row has pad at index 3
+    assert b["y_mask"][0, 1]
+    # BOS position of y_in is eos_id
+    assert (b["y_in"][:, 0] == cfg.eos_id).all()
+
+
+def test_flagged_composites_are_composite():
+    # sanity: the diagnostic set in eval.py really is composite and needs divisors > 7
+    for n in (121, 143, 169, 187):
+        assert any(n % d == 0 for d in range(2, int(n ** 0.5) + 1))
+        assert all(n % d != 0 for d in (2, 3, 5, 7))
diff --git a/tests/test_models.py b/tests/test_models.py
new file mode 100644
index 0000000..7994e86
--- /dev/null
+++ b/tests/test_models.py
@@ -0,0 +1,77 @@
+import torch
+import torch.nn.functional as F
+
+from src.config import Config
+from src.data import build_examples, make_batch
+from src.model_api import build_model
+
+
+def _batch(cfg, inputs=(2, 7, 42, 99)):
+    ex = build_examples(list(inputs), cfg)
+    return make_batch(ex, cfg)
+
+
+def test_rnn_shapes_and_halting_range():
+    cfg = Config(model="rnn", max_steps=20, min_steps=2)
+    m = build_model(cfg)
+    b = _batch(cfg)
+    out = m(b["x"], b["y_in"])
+    assert out["logits"].shape == (4, b["y"].shape[1], cfg.vocab)
+    s = out["halt_steps"]
+    assert s.shape == (4,)
+    assert bool((s >= cfg.min_steps).all())
+    assert bool((s <= cfg.max_steps).all())
+
+
+def test_rnn_fixed_k_ablation():
+    cfg = Config(model="rnn", halting=False, max_steps=20)
+    m = build_model(cfg)
+    b = _batch(cfg)
+    out = m(b["x"], b["y_in"])
+    assert torch.allclose(out["halt_steps"], torch.full((4,), 20.0))
+
+
+def test_transformer_shapes():
+    cfg = Config(model="transformer")
+    m = build_model(cfg)
+    b = _batch(cfg)
+    out = m(b["x"], b["y_in"])
+    assert out["logits"].shape == (4, b["y"].shape[1], cfg.vocab)
+    assert out["halt_steps"] is None
+
+
+def test_transformer_integers_mode():
+    cfg = Config(model="transformer", vocab_mode="integers")
+    m = build_model(cfg)
+    b = _batch(cfg)
+    out = m(b["x"], b["y_in"])
+    assert out["logits"].shape == (4, b["y"].shape[1], cfg.vocab)
+
+
+def test_param_counts_logged_not_gated():
+    r = build_model(Config(model="rnn")).param_count()
+    t = build_model(Config(model="transformer")).param_count()
+    assert r > 1000 and t > 1000
+
+
+def test_forward_backward_no_nan_both_models():
+    for model_name in ("rnn", "transformer"):
+        cfg = Config(model=model_name)
+        m = build_model(cfg)
+        b = _batch(cfg)
+        out = m(b["x"], b["y_in"])
+        y_safe = b["y"].clamp(max=cfg.vocab - 1)
+        loss = F.cross_entropy(out["logits"].reshape(-1, cfg.vocab), y_safe.reshape(-1), reduction="none")
+        mask = b["y_mask"].float().reshape(-1)
+        loss = (loss * mask).sum() / mask.sum()
+        loss.backward()
+        assert torch.isfinite(loss).item(), f"{model_name} loss NaN"
+
+
+def test_greedy_decode_shape():
+    from src.model_api import greedy_decode
+    cfg = Config(model="rnn")
+    m = build_model(cfg)
+    x = torch.tensor([[4, 2], [9, 9]], dtype=torch.long)
+    gen = greedy_decode(m, x, cfg)
+    assert gen.shape == (2, cfg.max_out_len)