Add files using upload-large-folder tool

code/notebooks/0101-autoregressive-likelihood-sft.ipynb

@@ -0,0 +1,309 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 0101 - Autoregressive LM likelihood, token vs sequence view, SFT = MLE = forward KL\n",
+    "Companion to `notes/0101 - Autoregressive Likelihood and SFT.pdf`.\n",
+    "Everything is fp64.  The model is a tiny autoregressive policy whose whole sequence space (121 sequences) is enumerated,\n",
+    "so each identity is checked EXACTLY (error ~1e-16), not by sampling.  Pattern: naive loop first, vectorised second, assert they agree."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "source": [
+    "import sys, math, itertools\n",
+    "sys.path.insert(0, \"..\")\n",
+    "from _m1_lib import *\n",
+    "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n",
+    "theta_s = lm.theta0.clone()                       # student\n",
+    "theta_T = TinyLM(V=4, L=4, seed=7, scale=1.5).theta0   # teacher p (same family -> realizable)\n",
+    "print(\"sequences:\", lm.N, \" params:\", theta_s.numel())"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "sequences: 121  params: 36\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. The model is a proper distribution on sequences:  sum_o pi(o) = 1   (note, Sec. 3.1)\n",
+    "Naive: multiply conditional probabilities one token at a time with a python loop.  Efficient: gather from one batched log_softmax."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "source": [
+    "p_naive = torch.stack([lm.seq_logprob_naive(theta_s, s).exp() for s in lm.seqs])\n",
+    "p_vec = lm.probs(theta_s)\n",
+    "check(\"naive == vectorised sequence probability\", maxabs(p_naive, p_vec), 1e-15)\n",
+    "check(\"sum_o pi(o) = 1 (normalisation, truncated at L)\", abs(float(p_vec.sum()) - 1.0), 1e-14)\n",
+    "# per-state conditional distributions each sum to 1 (this is what the induction in the note uses)\n",
+    "worst = max(abs(float(lm.cond_probs(theta_s, s).sum()) - 1) for s in lm.prefixes())\n",
+    "check(\"every conditional pi(.|o_<t) sums to 1\", worst, 1e-15)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|naive == vectorised sequence probability|0.000e+00|tol=1.0e-15|PASS\n",
+      "CHECK|exact|sum_o pi(o) = 1 (normalisation, truncated at L)|1.110e-16|tol=1.0e-14|PASS\n",
+      "CHECK|exact|every conditional pi(.|o_<t) sums to 1|2.220e-16|tol=1.0e-15|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Forward KL:  CE(p, pi) = H(p) + KL(p || pi)     (note, Eq. for the SFT loss)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "source": [
+    "p_T = lm.probs(theta_T)\n",
+    "lp_s = lm.seq_logprob(theta_s)\n",
+    "CE = -(p_T * lp_s).sum()\n",
+    "H = -(p_T * p_T.log()).sum()\n",
+    "KLf = (p_T * (p_T.log() - lp_s)).sum()\n",
+    "check(\"CE = H + KL(p||pi)\", abs(float(CE - H - KLf)), 1e-14)\n",
+    "check(\"KL(p||pi) >= 0 (Gibbs): min(0, KL)\", abs(min(0.0, float(KLf))), 0.0)\n",
+    "# reverse KL differs (what RL optimises is the *reverse* direction, see 0109)\n",
+    "KLr = (p_vec * (lp_s - p_T.log())).sum()\n",
+    "print(f\"forward KL(p||pi)={float(KLf):.6f}   reverse KL(pi||p)={float(KLr):.6f}   (different objects)\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|CE = H + KL(p||pi)|0.000e+00|tol=1.0e-14|PASS\n",
+      "CHECK|exact|KL(p||pi) >= 0 (Gibbs): min(0, KL)|0.000e+00|tol=0.0e+00|PASS\n",
+      "forward KL(p||pi)=1.221839   reverse KL(pi||p)=1.094181   (different objects)\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Chain rule of KL: sequence KL = sum over positions of the *expected conditional* KL  (token-level view)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "source": [
+    "def prefix_prob(theta, prefix):\n",
+    "    pr = 1.0\n",
+    "    for k, a in enumerate(prefix):\n",
+    "        pr = pr * lm.cond_probs(theta, prefix[:k])[a]\n",
+    "    return pr\n",
+    "tok_KL = 0.0\n",
+    "for s in lm.prefixes():\n",
+    "    w = prefix_prob(theta_T, s)                        # P_p(o_<t = s)   (only non-EOS prefixes are non-terminal states)\n",
+    "    cp, cq = lm.cond_probs(theta_T, s), lm.cond_probs(theta_s, s)\n",
+    "    tok_KL = tok_KL + w * (cp * (cp.log() - cq.log())).sum()\n",
+    "check(\"KL(p||pi)_sequence = sum_t E_p[KL_t]  (chain rule)\", abs(float(KLf - tok_KL)), 1e-14)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|KL(p||pi)_sequence = sum_t E_p[KL_t]  (chain rule)|2.220e-16|tol=1.0e-14|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Hand-derived gradient of the SFT loss vs autograd\n",
+    "For a softmax row, d log pi(a)/d logit_b = 1[b=a] - pi(b).  Our logits are tok[prev]+pos[t], so each token adds that vector to the\n",
+    "tok[prev] slot and to the pos[t] slot.  Naive: python loops over sequences and positions."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "source": [
+    "V, L = lm.V, lm.L\n",
+    "def hand_score(theta, seq):\n",
+    "    \"\"\"grad_theta log pi(seq), derived by hand: sum_t (e_{a_t} - pi(.|ctx_t)) scattered into the tok / pos slots.\"\"\"\n",
+    "    g = torch.zeros_like(theta)\n",
+    "    prev = V\n",
+    "    for t, a in enumerate(seq):\n",
+    "        pi = lm.cond_probs(theta, seq[:t])\n",
+    "        d = -pi.clone(); d[a] += 1.0\n",
+    "        g[prev * V:(prev + 1) * V] += d\n",
+    "        g[lm.n_tok + t * V: lm.n_tok + (t + 1) * V] += d\n",
+    "        prev = a\n",
+    "    return g\n",
+    "th = theta_s.clone().requires_grad_(True)\n",
+    "S_auto = score_matrix(lm, theta_s)\n",
+    "S_hand = torch.stack([hand_score(theta_s, s) for s in lm.seqs])\n",
+    "check(\"hand-derived score == autograd score (all 121 sequences)\", maxabs(S_auto, S_hand), 1e-14)\n",
+    "grad_CE_hand = -(p_T[:, None] * S_hand).sum(0)            # d/dtheta CE = -E_p[score]\n",
+    "th = theta_s.clone().requires_grad_(True)\n",
+    "grad_CE_auto = torch.autograd.grad(-(p_T * lm.seq_logprob(th)).sum(), th)[0]\n",
+    "check(\"grad CE: -E_p[score] == autograd\", maxabs(grad_CE_hand, grad_CE_auto), 1e-14)\n",
+    "# at the teacher the MLE gradient vanishes (p is realisable)\n",
+    "S_T = score_matrix(lm, theta_T)\n",
+    "check(\"grad CE at teacher = 0\", float((p_T[:, None] * S_T).sum(0).abs().max()), 1e-14)\n",
+    "# E_pi[score] = 0 for every theta (the identity used in 0102)\n",
+    "check(\"E_pi[score] = 0 (zero-mean score)\", float((p_vec[:, None] * S_auto).sum(0).abs().max()), 1e-14)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|hand-derived score == autograd score (all 121 sequences)|1.110e-16|tol=1.0e-14|PASS\n",
+      "CHECK|exact|grad CE: -E_p[score] == autograd|9.714e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|grad CE at teacher = 0|4.857e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|E_pi[score] = 0 (zero-mean score)|5.551e-17|tol=1.0e-14|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Loss normalisations are NOT equivalent (token-sum vs sequence-mean vs global token-mean)\n",
+    "SFT gradient = -E_p[sum_t grad log pi_t].  Dividing each sequence by its own length (seq-mean) re-weights tokens by 1/|o|."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "source": [
+    "Ttens = token_score_tensor(lm, theta_s)                  # [N, L, P]\n",
+    "len_ = lm.lengths\n",
+    "g_sum = -(p_T[:, None] * Ttens.sum(1)).sum(0)\n",
+    "g_seq = -(p_T[:, None] * (Ttens.sum(1) / len_[:, None])).sum(0)\n",
+    "Lmax = float(L)\n",
+    "g_const = -(p_T[:, None] * (Ttens.sum(1) / Lmax)).sum(0)\n",
+    "check(\"token-sum gradient == MLE gradient\", maxabs(g_sum, grad_CE_hand), 1e-14)\n",
+    "check(\"constant-denominator loss == MLE gradient / Lmax\", maxabs(g_const * Lmax, g_sum), 1e-14)\n",
+    "rel = float((g_seq - g_sum / float(p_T @ len_)).norm() / g_sum.norm())\n",
+    "print(f\"seq-mean gradient vs MLE gradient: relative deviation {rel:.3f} (NOT a pure rescale -> reweights tokens by 1/|o|)\")\n",
+    "assert rel > 1e-3"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|token-sum gradient == MLE gradient|9.714e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|constant-denominator loss == MLE gradient / Lmax|0.000e+00|tol=1.0e-14|PASS\n",
+      "seq-mean gradient vs MLE gradient: relative deviation 0.444 (NOT a pure rescale -> reweights tokens by 1/|o|)\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Sampling check (MC): empirical-risk gradient converges to the population one"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "source": [
+    "gen = torch.Generator().manual_seed(123)\n",
+    "n = 400_000\n",
+    "idx = torch.multinomial(p_T, n, replacement=True, generator=gen)\n",
+    "G_emp = -S_auto[idx].mean(0)\n",
+    "se = (S_auto[idx].var(0) / n).sqrt()\n",
+    "ok_ = se > 0\n",
+    "z = float(((G_emp - grad_CE_hand).abs()[ok_] / se[ok_]).max())\n",
+    "assert float((G_emp - grad_CE_hand).abs()[~ok_].max()) < 1e-14   # unused parameters have identically zero gradient\n",
+    "check(\"MC SFT gradient within 4 SE (max z over components)\", z, 4.0, kind=\"mc\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|mc|MC SFT gradient within 4 SE (max z over components)|1.761e+00|tol=4.0e+00|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7. SFT = maximum likelihood = forward KL minimisation: L-BFGS from a random init recovers the teacher"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "source": [
+    "gen_i = torch.Generator().manual_seed(77)                       # seeded: the quoted numbers must reproduce\n",
+    "th = (theta_s + 0.5 * torch.randn(theta_s.numel(), generator=gen_i)).requires_grad_(True)\n",
+    "opt = torch.optim.LBFGS([th], lr=1.0, max_iter=500, tolerance_grad=1e-14, tolerance_change=1e-16, line_search_fn=\"strong_wolfe\")\n",
+    "def closure():\n",
+    "    opt.zero_grad(); loss = -(p_T * lm.seq_logprob(th)).sum(); loss.backward(); return loss\n",
+    "for _ in range(6): opt.step(closure)\n",
+    "KLend = float((p_T * (p_T.log() - lm.seq_logprob(th.detach()))).sum())\n",
+    "check(\"KL(p||pi) after MLE (should be ~0)\", KLend, 1e-10)\n",
+    "check(\"recovered sequence distribution == teacher\", maxabs(lm.probs(th.detach()), p_T), 1e-6)\n",
+    "print(\"done; checks:\", len(CHECKS))"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|KL(p||pi) after MLE (should be ~0)|4.753e-17|tol=1.0e-10|PASS\n",
+      "CHECK|exact|recovered sequence distribution == teacher|2.203e-09|tol=1.0e-06|PASS\n",
+      "done; checks: 15\n"
+     ]
+    }
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

code/notebooks/0102-policy-gradient-baselines.ipynb

@@ -0,0 +1,347 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 0102 - Policy gradient from first principles, REINFORCE, baselines\n",
+    "Companion to `notes/0102 - Policy Gradient and Baselines.pdf`.  Exact enumeration of the 121-sequence toy LM; fp64."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "source": [
+    "import sys, itertools\n",
+    "sys.path.insert(0, \"..\")\n",
+    "from _m1_lib import *\n",
+    "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n",
+    "theta = lm.theta0.clone()\n",
+    "R = default_reward(lm)                       # binary verifier reward\n",
+    "p = lm.probs(theta)\n",
+    "S = score_matrix(lm, theta)                  # [N, P]  grad log pi(o)\n",
+    "P_ = theta.numel()\n",
+    "print(\"N =\", lm.N, \" P =\", P_, \" E[R] =\", float(p @ R))"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "N = 121  P = 36  E[R] = 0.16592597945219617\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Log-derivative trick:  grad J = E_pi[R * score]   (note Eq. 2)\n",
+    "J(theta) = sum_o pi_theta(o) R(o).  Three independent computations of its gradient."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "source": [
+    "th = theta.clone().requires_grad_(True)\n",
+    "J = (lm.probs(th) * R).sum()\n",
+    "g_auto = torch.autograd.grad(J, th)[0]                       # differentiate the *definition* (through pi)\n",
+    "g_ld = (p[:, None] * R[:, None] * S).sum(0)                   # E_pi[R * score]  (the derived formula)\n",
+    "def Jfun(t): return float((lm.probs(t) * R).sum())\n",
+    "h = 1e-6\n",
+    "g_fd = torch.zeros_like(theta)\n",
+    "for i in range(P_):                                           # naive central finite differences\n",
+    "    e = torch.zeros_like(theta); e[i] = h\n",
+    "    g_fd[i] = (Jfun(theta + e) - Jfun(theta - e)) / (2 * h)\n",
+    "check(\"log-derivative formula == autograd of J\", maxabs(g_ld, g_auto), 1e-14)\n",
+    "check(\"log-derivative formula == central finite differences\", maxabs(g_ld, g_fd), 1e-8)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|log-derivative formula == autograd of J|1.388e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|log-derivative formula == central finite differences|3.363e-11|tol=1.0e-08|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Baselines are free:  E[b * score] = 0  for any b that does not depend on the sampled response"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "source": [
+    "check(\"E_pi[score] = 0\", float((p[:, None] * S).sum(0).abs().max()), 1e-14)\n",
+    "for b in [0.0, float(p @ R), 5.0, -3.0]:\n",
+    "    g_b = (p[:, None] * (R - b)[:, None] * S).sum(0)\n",
+    "    check(f\"unbiased with constant baseline b={b:+.3f}\", maxabs(g_b, g_auto), 1e-13)\n",
+    "# state-dependent baseline b_t(o_<t): arbitrary function of the prefix  (note Eq. for token-level baseline)\n",
+    "T = token_score_tensor(lm, theta)                            # [N, L, P]  token scores\n",
+    "gen = torch.Generator().manual_seed(5)\n",
+    "tab = {}\n",
+    "def bpref(prefix):\n",
+    "    if prefix not in tab: tab[prefix] = float(torch.randn(1, generator=gen)) * 3.0\n",
+    "    return tab[prefix]\n",
+    "B = torch.zeros(lm.N, lm.L)\n",
+    "for i, s in enumerate(lm.seqs):\n",
+    "    for t in range(len(s)): B[i, t] = bpref(s[:t])\n",
+    "g_state = (p[:, None] * ((R[:, None] - B)[:, :, None] * T).sum(1)).sum(0)\n",
+    "check(\"unbiased with arbitrary state-dependent baseline b(o_<t)\", maxabs(g_state, g_auto), 1e-13)\n",
+    "# a baseline that depends on the token being scored breaks the proof (counter-example)\n",
+    "B_bad = 0.7 * (lm.tokens == 1).double() * lm.mask\n",
+    "g_bad = (p[:, None] * ((R[:, None] - B_bad)[:, :, None] * T).sum(1)).sum(0)\n",
+    "print(\"baseline depending on o_t: bias norm =\", float((g_bad - g_auto).norm()))\n",
+    "assert float((g_bad - g_auto).norm()) > 1e-3\n",
+    "# and the extreme case b = R(o) (the sample's own reward) kills the gradient completely\n",
+    "g_self = (p[:, None] * (R - R)[:, None] * S).sum(0)\n",
+    "print(\"self-baseline gradient norm:\", float(g_self.norm()), \" true grad norm:\", float(g_auto.norm()))"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|E_pi[score] = 0|5.551e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|unbiased with constant baseline b=+0.000|1.388e-17|tol=1.0e-13|PASS\n",
+      "CHECK|exact|unbiased with constant baseline b=+0.166|2.776e-17|tol=1.0e-13|PASS\n",
+      "CHECK|exact|unbiased with constant baseline b=+5.000|2.078e-16|tol=1.0e-13|PASS\n",
+      "CHECK|exact|unbiased with constant baseline b=-3.000|2.137e-16|tol=1.0e-13|PASS\n",
+      "CHECK|exact|unbiased with arbitrary state-dependent baseline b(o_<t)|1.110e-16|tol=1.0e-13|PASS\n",
+      "baseline depending on o_t: bias norm = 0.1657102369175368\n",
+      "self-baseline gradient norm: 0.0  true grad norm: 0.18246238486098645\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Variance: exact quadratic in b, optimal baseline, and the \"(0, 2 b*)\" interval\n",
+    "Var(b) = E||(R-b) s||^2 - ||grad J||^2  with  E||(R-b)s||^2 = E[R^2||s||^2] - 2b E[R||s||^2] + b^2 E||s||^2."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "source": [
+    "s2 = (S ** 2).sum(1)                                   # ||s||^2 per sequence\n",
+    "A0, A1, A2 = float((p * R ** 2 * s2).sum()), float((p * R * s2).sum()), float((p * s2).sum())\n",
+    "bstar = A1 / A2\n",
+    "def var_direct(b):\n",
+    "    g = (R - b)[:, None] * S                            # per-sequence estimator value\n",
+    "    return float((p * (g ** 2).sum(1)).sum() - (p[:, None] * g).sum(0).norm() ** 2)\n",
+    "def var_formula(b): return A0 - 2 * b * A1 + b * b * A2 - float(g_auto.norm() ** 2)\n",
+    "for b in [0.0, 0.3, 1.0, 2.0]:\n",
+    "    check(f\"variance formula at b={b}\", abs(var_direct(b) - var_formula(b)), 1e-13)\n",
+    "grid = torch.linspace(-2, 4, 6001)\n",
+    "vs = torch.tensor([var_formula(float(b)) for b in grid])\n",
+    "check(\"b* = E[R||s||^2]/E||s||^2 minimises variance (grid min vs V(b*))\", max(0.0, var_formula(bstar) - float(vs.min())), 1e-12)\n",
+    "check(\"minimiser on grid is within grid step of b*\", abs(float(grid[vs.argmin()]) - bstar), 2e-3)\n",
+    "EvR = float(p @ R)\n",
+    "print(f\"b*={bstar:.4f}  E[R]={EvR:.4f}   Var(0)={var_formula(0.0):.4f}  Var(E[R])={var_formula(EvR):.4f}  Var(b*)={var_formula(bstar):.4f}\")\n",
+    "# which side of E[R] does b* fall on?  b* - E[R] = Cov_pi(R, ||s||^2) / E||s||^2   (exact identity); sign = sign of the covariance\n",
+    "cov_Rs2 = float((p * R * s2).sum() - EvR * (p * s2).sum())\n",
+    "check(\"b* - E[R] == Cov_pi(R, ||s||^2) / E||s||^2  (exact identity)\", abs((bstar - EvR) - cov_Rs2 / A2), 1e-14)\n",
+    "m1_ = float((p * R * s2).sum() / (p * R).sum()); m0_ = float((p * (1 - R) * s2).sum() / (p * (1 - R)).sum())\n",
+    "print(f\"E[||s||^2 | R=1] = {m1_:.3f}   E[||s||^2 | R=0] = {m0_:.3f}   -> successes carry the larger score norm, so b* > E[R]\")\n",
+    "assert (m1_ > m0_) == (bstar > EvR)\n",
+    "# interval claim: any b in (0, 2b*) has Var(b) < Var(0); b outside is worse\n",
+    "inside = all(var_formula(float(b)) < var_formula(0.0) for b in torch.linspace(0.01, 2 * bstar - 0.01, 50))\n",
+    "outside = all(var_formula(float(b)) > var_formula(0.0) for b in [2 * bstar + 0.01, 2 * bstar + 1.0, -0.01, -1.0])\n",
+    "assert inside and outside\n",
+    "check(\"Var(b) < Var(0) iff b in (0, 2 b*)  (50 inside, 4 outside points)\", 0.0, 0.0)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|variance formula at b=0.0|0.000e+00|tol=1.0e-13|PASS\n",
+      "CHECK|exact|variance formula at b=0.3|1.110e-16|tol=1.0e-13|PASS\n",
+      "CHECK|exact|variance formula at b=1.0|2.220e-16|tol=1.0e-13|PASS\n",
+      "CHECK|exact|variance formula at b=2.0|8.882e-16|tol=1.0e-13|PASS\n",
+      "CHECK|exact|b* = E[R||s||^2]/E||s||^2 minimises variance (grid min vs V(b*))|0.000e+00|tol=1.0e-12|PASS\n",
+      "CHECK|exact|minimiser on grid is within grid step of b*|1.020e-04|tol=2.0e-03|PASS\n",
+      "b*=0.3019  E[R]=0.1659   Var(0)=0.6898  Var(E[R])=0.5158  Var(b*)=0.4715\n",
+      "CHECK|exact|b* - E[R] == Cov_pi(R, ||s||^2) / E||s||^2  (exact identity)|0.000e+00|tol=1.0e-14|PASS\n",
+      "E[||s||^2 | R=1] = 4.358   E[||s||^2 | R=0] = 2.005   -> successes carry the larger score norm, so b* > E[R]\n",
+      "CHECK|exact|Var(b) < Var(0) iff b in (0, 2 b*)  (50 inside, 4 outside points)|0.000e+00|tol=0.0e+00|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Naive vs efficient estimator (Monte-Carlo), and agreement with the exact gradient\n",
+    "naive: one backward per sample, python loop.  efficient: ONE backward through the surrogate  -mean(sg(A) * log pi)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "source": [
+    "n = 2000\n",
+    "gen = torch.Generator().manual_seed(11)\n",
+    "idx = torch.multinomial(p, n, replacement=True, generator=gen)\n",
+    "b = EvR\n",
+    "# naive\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "acc = torch.zeros_like(theta)\n",
+    "for i in idx.tolist():\n",
+    "    lp = lm.seq_logprob(th)[i]\n",
+    "    acc += (float(R[i]) - b) * torch.autograd.grad(lp, th)[0]\n",
+    "g_naive = acc / n\n",
+    "# efficient\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "surrogate = ((R[idx] - b).detach() * lm.seq_logprob(th)[idx]).mean()    # gradient of this = REINFORCE estimate\n",
+    "g_eff = torch.autograd.grad(surrogate, th)[0]\n",
+    "check(\"naive per-sample loop == single backward through surrogate\", maxabs(g_naive, g_eff), 1e-13)\n",
+    "# MC accuracy against the exact gradient\n",
+    "n2 = 400_000\n",
+    "idx2 = torch.multinomial(p, n2, replacement=True, generator=gen)\n",
+    "est = ((R[idx2] - b)[:, None] * S[idx2])\n",
+    "g_mc, se = est.mean(0), est.std(0) / n2 ** 0.5\n",
+    "ok_ = se > 0\n",
+    "z = float(((g_mc - g_auto).abs()[ok_] / se[ok_]).max())\n",
+    "check(\"MC REINFORCE gradient within 4 SE of exact (max z)\", z, 4.0, kind=\"mc\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|naive per-sample loop == single backward through surrogate|8.604e-16|tol=1.0e-13|PASS\n",
+      "CHECK|mc|MC REINFORCE gradient within 4 SE of exact (max z)|2.001e+00|tol=4.0e+00|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Group baselines: leave-one-out is unbiased, \"mean including self\" is biased by (G-1)/G\n",
+    "Exact enumeration of ALL G-tuples of samples (G=2: 14641 tuples, G=3: 1.77M tuples, evaluated per first sample)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "source": [
+    "N = lm.N\n",
+    "RR = R[None, :] ; pp = p[None, :]\n",
+    "def loo_and_mean(G):\n",
+    "    \"\"\"exact E[(R_1 - baseline_1) s_1] for baseline = mean of others (LOO) and mean of all (incl. self).\"\"\"\n",
+    "    # sum over samples 2..G of rewards: build the joint over (o_2..o_G) as a tensor of (reward sum, prob)\n",
+    "    if G == 2:\n",
+    "        rs = R[None, :].expand(N, N); w = p[None, :].expand(N, N)\n",
+    "        sum_others_R = rs; prob_others = w\n",
+    "    elif G == 3:\n",
+    "        sum_others_R = (R[:, None] + R[None, :]).reshape(-1)[None, :].expand(N, -1)\n",
+    "        prob_others = (p[:, None] * p[None, :]).reshape(-1)[None, :].expand(N, -1)\n",
+    "        # N x N^2 : first sample index x joint of other two\n",
+    "    Ro1 = R[:, None]\n",
+    "    base_loo = sum_others_R / (G - 1)\n",
+    "    base_all = (Ro1 + sum_others_R) / G\n",
+    "    wloo = (prob_others * (Ro1 - base_loo)).sum(1)       # E over others of (R1 - b)\n",
+    "    wall = (prob_others * (Ro1 - base_all)).sum(1)\n",
+    "    g_loo = (p * wloo)[:, None] * S\n",
+    "    g_all = (p * wall)[:, None] * S\n",
+    "    return g_loo.sum(0), g_all.sum(0)\n",
+    "for G in (2, 3):\n",
+    "    gl, ga = loo_and_mean(G)\n",
+    "    check(f\"G={G}: leave-one-out baseline is exactly unbiased\", maxabs(gl, g_auto), 1e-13)\n",
+    "    check(f\"G={G}: mean-including-self gives exactly (G-1)/G * grad J\", maxabs(ga, (G - 1) / G * g_auto), 1e-13)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|G=2: leave-one-out baseline is exactly unbiased|2.082e-17|tol=1.0e-13|PASS\n",
+      "CHECK|exact|G=2: mean-including-self gives exactly (G-1)/G * grad J|1.041e-17|tol=1.0e-13|PASS\n",
+      "CHECK|exact|G=3: leave-one-out baseline is exactly unbiased|2.082e-17|tol=1.0e-13|PASS\n",
+      "CHECK|exact|G=3: mean-including-self gives exactly (G-1)/G * grad J|6.939e-18|tol=1.0e-13|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Causality / reward-to-go (per-token rewards): both unbiased, reward-to-go has lower variance here\n",
+    "Per-token reward r_t = 1[o_t == 1].  Full return G = sum_t r_t;  reward-to-go G_t = sum_{t'>=t} r_t'."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "source": [
+    "rt = (lm.tokens == 1).double() * lm.mask                    # [N, L]\n",
+    "Gfull = rt.sum(1)\n",
+    "Gto = rt.flip(1).cumsum(1).flip(1) * lm.mask                # reward-to-go\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "gJ = torch.autograd.grad((lm.probs(th) * Gfull).sum(), th)[0]\n",
+    "est_full = (T.sum(1) * Gfull[:, None])                       # sum_t score_t * G\n",
+    "est_rtg = (T * Gto[:, :, None]).sum(1)                       # sum_t score_t * G_t\n",
+    "g_full = (p[:, None] * est_full).sum(0); g_rtg = (p[:, None] * est_rtg).sum(0)\n",
+    "check(\"full-return estimator unbiased\", maxabs(g_full, gJ), 1e-13)\n",
+    "check(\"reward-to-go estimator unbiased (dropped terms E[score_t r_t'] , t'<t, vanish)\", maxabs(g_rtg, gJ), 1e-13)\n",
+    "v_full = float((p * (est_full ** 2).sum(1)).sum() - gJ.norm() ** 2)\n",
+    "v_rtg = float((p * (est_rtg ** 2).sum(1)).sum() - gJ.norm() ** 2)\n",
+    "print(f\"trace-variance full return {v_full:.4f}   reward-to-go {v_rtg:.4f}\")\n",
+    "assert v_rtg < v_full\n",
+    "# for an OUTCOME reward (one number at the end) reward-to-go == full return: nothing to drop\n",
+    "Gto_out = R[:, None] * lm.mask\n",
+    "check(\"outcome reward: reward-to-go == outcome at every t (no causality gain)\", maxabs(Gto_out.sum(1), R * lm.lengths), 1e-15)\n",
+    "print(\"done; checks:\", len(CHECKS))"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|full-return estimator unbiased|4.163e-17|tol=1.0e-13|PASS\n",
+      "CHECK|exact|reward-to-go estimator unbiased (dropped terms E[score_t r_t'] , t'<t, vanish)|4.163e-17|tol=1.0e-13|PASS\n",
+      "trace-variance full return 2.8551   reward-to-go 1.6247\n",
+      "CHECK|exact|outcome reward: reward-to-go == outcome at every t (no causality gain)|0.000e+00|tol=1.0e-15|PASS\n",
+      "done; checks: 25\n"
+     ]
+    }
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

code/notebooks/0103-importance-sampling-offpolicy.ipynb

@@ -0,0 +1,357 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 0103 - Importance sampling for off-policy updates; ratio variance; the product-of-token-ratios explosion\n",
+    "Companion to `notes/0103 - Importance Sampling and Off-Policy Updates.pdf`.  Exact enumeration, fp64."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "source": [
+    "import sys, math\n",
+    "sys.path.insert(0, \"..\")\n",
+    "from _m1_lib import *"
+   ],
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. IS identity, unit mean, and Var_q[w] = chi^2 = E_p[w] - 1   (note Sec. 2.1)\n",
+    "q = pi_old (sampler), p = pi_theta (target).  Sequence-level weight w(o) = p(o)/q(o)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "source": [
+    "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n",
+    "th_old = lm.theta0.clone()\n",
+    "th_new = perturb(th_old, 0.5, seed=3)\n",
+    "q, p = lm.probs(th_old), lm.probs(th_new)\n",
+    "w = p / q\n",
+    "f = default_reward(lm) + 0.3 * lm.lengths                    # an arbitrary test function f(o)\n",
+    "check(\"E_q[w f] = E_p[f]  (IS identity)\", abs(float((q * w * f).sum() - (p * f).sum())), 1e-14)\n",
+    "check(\"E_q[w] = 1\", abs(float((q * w).sum()) - 1.0), 1e-14)\n",
+    "chi2 = float((p * w).sum()) - 1.0\n",
+    "check(\"Var_q[w] = E_p[w] - 1 = chi^2(p||q)\", abs(float((q * w ** 2).sum() - 1.0) - chi2), 1e-14)\n",
+    "# variance of the IS estimator of E_p[f]:  Var_q[w f] = E_p[w f^2] - (E_p f)^2\n",
+    "v_direct = float((q * (w * f) ** 2).sum() - (q * w * f).sum() ** 2)\n",
+    "v_formula = float((p * w * f ** 2).sum() - (p * f).sum() ** 2)\n",
+    "check(\"Var_q[w f] = E_p[w f^2] - (E_p f)^2\", abs(v_direct - v_formula), 1e-13)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|E_q[w f] = E_p[f]  (IS identity)|0.000e+00|tol=1.0e-14|PASS\n",
+      "CHECK|exact|E_q[w] = 1|0.000e+00|tol=1.0e-14|PASS\n",
+      "CHECK|exact|Var_q[w] = E_p[w] - 1 = chi^2(p||q)|0.000e+00|tol=1.0e-14|PASS\n",
+      "CHECK|exact|Var_q[w f] = E_p[w f^2] - (E_p f)^2|0.000e+00|tol=1.0e-13|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Self-normalised IS: bias O(1/n) and the effective sample size  ESS = (sum w)^2 / sum w^2  ~  n / (1 + chi^2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "source": [
+    "gen = torch.Generator().manual_seed(21)\n",
+    "def snis(n, reps):\n",
+    "    est, ess = [], []\n",
+    "    for _ in range(reps):\n",
+    "        idx = torch.multinomial(q, n, replacement=True, generator=gen)\n",
+    "        wi = w[idx]\n",
+    "        est.append(float((wi * f[idx]).sum() / wi.sum()))\n",
+    "        ess.append(float(wi.sum() ** 2 / (wi ** 2).sum()))\n",
+    "    return torch.tensor(est), torch.tensor(ess)\n",
+    "target = float((p * f).sum())\n",
+    "for n in (4, 16, 64, 256):\n",
+    "    est, ess = snis(n, 40000 // (1 if n <= 16 else 4))\n",
+    "    bias = float(est.mean() - target); se = float(est.std() / len(est) ** 0.5)\n",
+    "    print(f\"n={n:4d}  SNIS bias={bias:+.5f} (se {se:.5f})   mean ESS={float(ess.mean()):.2f}   n/(1+chi2)={n/(1+chi2):.2f}\")\n",
+    "est, ess = snis(256, 4000)\n",
+    "check(\"MC: ESS at n=256 within 15% of n/(1+chi^2)\", abs(float(ess.mean()) / (256 / (1 + chi2)) - 1), 0.15, kind=\"mc\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "n=   4  SNIS bias=-0.01069 (se 0.00161)   mean ESS=3.96   n/(1+chi2)=3.94\n",
+      "n=  16  SNIS bias=-0.00276 (se 0.00082)   mean ESS=15.78   n/(1+chi2)=15.76\n",
+      "n=  64  SNIS bias=-0.00104 (se 0.00082)   mean ESS=63.07   n/(1+chi2)=63.04\n",
+      "n= 256  SNIS bias=-0.00018 (se 0.00040)   mean ESS=252.20   n/(1+chi2)=252.17\n",
+      "CHECK|mc|MC: ESS at n=256 within 15% of n/(1+chi^2)|7.446e-05|tol=1.5e-01|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Product of token ratios: exact chi^2 growth with horizon T  (long responses)\n",
+    "Sequence weight over the first T tokens: w_T = prod_{t<=T} r_t.  E_q[w_T] = 1 for every T but E_q[w_T^2] grows geometrically."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "source": [
+    "lmL = TinyLM(V=3, L=12, seed=4, scale=1.0)\n",
+    "tok = lmL.theta0[:lmL.n_tok].reshape(lmL.V + 1, lmL.V); tok[:, lmL.EOS] -= 2.5      # make EOS rare -> long sequences\n",
+    "thq = lmL.theta0.clone(); thp = perturb(thq, 0.6, seed=9)\n",
+    "lq, lp = lmL.token_logprobs(thq), lmL.token_logprobs(thp)\n",
+    "qL = lq.sum(1).exp()\n",
+    "check(\"long-horizon model: sum_o q(o) = 1\", abs(float(qL.sum()) - 1), 1e-12)\n",
+    "print(\"N sequences:\", lmL.N, \" mean length under q:\", float((qL * lmL.lengths).sum()))\n",
+    "logw_cum = (lp - lq).cumsum(1)                                # log w_T for T = 1..L  (masked tokens contribute 0 -> ratio 1 after EOS)\n",
+    "rows = []\n",
+    "for T in range(1, lmL.L + 1):\n",
+    "    wT = logw_cum[:, T - 1].exp()\n",
+    "    m1 = float((qL * wT).sum()); m2 = float((qL * wT ** 2).sum())\n",
+    "    med = float((qL * (wT < 1).double()).sum())               # P_q(w_T < 1): the typical sample has weight < 1\n",
+    "    rows.append((T, m1, m2 - 1, med, float((qL * logw_cum[:, T - 1]).sum())))\n",
+    "    check(f\"E_q[w_T] = 1 at T={T}\", abs(m1 - 1), 1e-12) if T in (1, 6, 12) else None\n",
+    "print(\" T   E[w]    chi2_T      P(w<1)   E[log w]\")\n",
+    "for r in rows: print(f\"{r[0]:2d}  {r[1]:.4f}  {r[2]:9.4f}   {r[3]:.4f}   {r[4]:+.4f}\")\n",
+    "chi_T = torch.tensor([r[2] for r in rows])\n",
+    "assert bool((chi_T[1:] >= chi_T[:-1] - 1e-12).all())\n",
+    "check(\"chi^2_T non-decreasing in T (E_q[w_T^2] only grows)\", 0.0, 0.0)\n",
+    "ratio_growth = float(chi_T[-1] / chi_T[0])\n",
+    "print(f\"chi2_12 / chi2_1 = {ratio_growth:.1f}\")\n",
+    "# chain rule:  E_q[log w_T] = - sum_{t<=T} E_q[KL_t(q || p)]   computed independently on the prefix tree\n",
+    "def tree_kl(thq, thp, T):\n",
+    "    tot = 0.0\n",
+    "    def rec(prefix, prob):\n",
+    "        nonlocal tot\n",
+    "        if len(prefix) >= T: return\n",
+    "        cq, cp = lmL.cond_probs(thq, prefix), lmL.cond_probs(thp, prefix)\n",
+    "        tot += prob * float((cq * (cq.log() - cp.log())).sum())\n",
+    "        for a in range(lmL.V - 1):\n",
+    "            rec(prefix + (a,), prob * float(cq[a]))\n",
+    "    rec((), 1.0)\n",
+    "    return tot\n",
+    "for T in (3, 8):\n",
+    "    check(f\"E_q[log w_T] = -sum_t E[KL_t] (tree recursion), T={T}\", abs(rows[T - 1][4] + tree_kl(thq, thp, T)), 1e-12)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|long-horizon model: sum_o q(o) = 1|2.220e-16|tol=1.0e-12|PASS\n",
+      "N sequences: 8191  mean length under q: 8.597261658540575\n",
+      "CHECK|exact|E_q[w_T] = 1 at T=1|2.220e-16|tol=1.0e-12|PASS\n",
+      "CHECK|exact|E_q[w_T] = 1 at T=6|0.000e+00|tol=1.0e-12|PASS\n",
+      "CHECK|exact|E_q[w_T] = 1 at T=12|0.000e+00|tol=1.0e-12|PASS\n",
+      " T   E[w]    chi2_T      P(w<1)   E[log w]\n",
+      " 1  1.0000     0.0040   0.8060   -0.0019\n",
+      " 2  1.0000     0.0042   0.8044   -0.0020\n",
+      " 3  1.0000     0.0333   0.2960   -0.0189\n",
+      " 4  1.0000     0.0373   0.4582   -0.0207\n",
+      " 5  1.0000     0.0446   0.4510   -0.0244\n",
+      " 6  1.0000     0.0457   0.4176   -0.0249\n",
+      " 7  1.0000     0.0487   0.4547   -0.0263\n",
+      " 8  1.0000     0.0500   0.4533   -0.0269\n",
+      " 9  1.0000     0.0511   0.4731   -0.0274\n",
+      "10  1.0000     0.0534   0.4954   -0.0285\n",
+      "11  1.0000     0.0535   0.4942   -0.0285\n",
+      "12  1.0000     0.0537   0.4936   -0.0286\n",
+      "CHECK|exact|chi^2_T non-decreasing in T (E_q[w_T^2] only grows)|0.000e+00|tol=0.0e+00|PASS\n",
+      "chi2_12 / chi2_1 = 13.3\n",
+      "CHECK|exact|E_q[log w_T] = -sum_t E[KL_t] (tree recursion), T=3|4.857e-17|tol=1.0e-12|PASS\n",
+      "CHECK|exact|E_q[log w_T] = -sum_t E[KL_t] (tree recursion), T=8|1.874e-16|tol=1.0e-12|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. iid-token closed form  E_q[w_T^2] = (1 + chi2_1)^T   and ESS collapse"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "source": [
+    "Vv = 3\n",
+    "qa = torch.tensor([0.5, 0.3, 0.2]); pa = torch.tensor([0.45, 0.33, 0.22])\n",
+    "chi1 = float((pa ** 2 / qa).sum()) - 1\n",
+    "def brute(T):\n",
+    "    seqs = torch.cartesian_prod(*[torch.arange(Vv)] * T) if T > 1 else torch.arange(Vv)[:, None]\n",
+    "    lqs = qa[seqs].log().sum(1); lps = pa[seqs].log().sum(1)\n",
+    "    return float((lqs.exp() * (lps - lqs).exp() ** 2).sum())\n",
+    "for T in (1, 3, 6, 9):\n",
+    "    check(f\"iid tokens: E_q[w_T^2] = (1+chi2_1)^T, T={T}\", abs(brute(T) - (1 + chi1) ** T) / (1 + chi1) ** T, 1e-12)\n",
+    "print(f\"per-token chi2_1 = {chi1:.5f}\")\n",
+    "for T in (1, 10, 100, 1000, 4000):\n",
+    "    print(f\"T={T:5d}   chi2_T = {(1+chi1)**T-1:12.3f}   ESS/n ~ 1/(1+chi2) = {(1+chi1)**(-T):.3e}\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|iid tokens: E_q[w_T^2] = (1+chi2_1)^T, T=1|0.000e+00|tol=1.0e-12|PASS\n",
+      "CHECK|exact|iid tokens: E_q[w_T^2] = (1+chi2_1)^T, T=3|4.310e-16|tol=1.0e-12|PASS\n",
+      "CHECK|exact|iid tokens: E_q[w_T^2] = (1+chi2_1)^T, T=6|0.000e+00|tol=1.0e-12|PASS\n",
+      "CHECK|exact|iid tokens: E_q[w_T^2] = (1+chi2_1)^T, T=9|4.060e-16|tol=1.0e-12|PASS\n",
+      "per-token chi2_1 = 0.01000\n",
+      "T=    1   chi2_T =        0.010   ESS/n ~ 1/(1+chi2) = 9.901e-01\n",
+      "T=   10   chi2_T =        0.105   ESS/n ~ 1/(1+chi2) = 9.053e-01\n",
+      "T=  100   chi2_T =        1.705   ESS/n ~ 1/(1+chi2) = 3.697e-01\n",
+      "T= 1000   chi2_T =    20958.156   ESS/n ~ 1/(1+chi2) = 4.771e-05\n",
+      "T= 4000   chi2_T = 192972369947321888.000   ESS/n ~ 1/(1+chi2) = 5.182e-18\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Naive vs stable weight computation (fp32 running product underflows; log-space is exact)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "source": [
+    "gen = torch.Generator().manual_seed(77)\n",
+    "def weights(T, sigma, dtype):\n",
+    "    logr = (-0.5 * sigma ** 2 + sigma * torch.randn(T, generator=gen)).to(torch.float64)\n",
+    "    naive = torch.tensor(1.0, dtype=dtype)\n",
+    "    for x in logr.exp().to(dtype): naive = naive * x          # naive running product of token ratios\n",
+    "    stable = logr.sum().exp()                                  # log-space\n",
+    "    return float(naive), float(stable), float(logr.sum())\n",
+    "n64, s64, _ = weights(50, 0.1, torch.float64)\n",
+    "check(\"T=50: naive fp64 product == exp(sum log r)\", abs(n64 - s64) / s64, 1e-12)\n",
+    "n32, s64b, lw = weights(4000, 0.5, torch.float32)\n",
+    "print(f\"T=4000 sigma=0.5:  log w = {lw:.1f};  naive fp32 product = {n32};  log-space fp64 = {s64b:.3e}\")\n",
+    "assert n32 == 0.0 and math.isfinite(lw)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|T=50: naive fp64 product == exp(sum log r)|0.000e+00|tol=1.0e-12|PASS\n",
+      "T=4000 sigma=0.5:  log w = -521.5;  naive fp32 product = 0.0;  log-space fp64 = 3.410e-227\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Token-level surrogate (PPO-style ratio) vs the exact off-policy gradient\n",
+    "Exact:   grad J(theta) = E_old[ w(o) (R-b) sum_t grad log pi_theta(o_t|s_t) ]  (sequence-level weight, all tokens).\n",
+    "Token-level surrogate:  L(theta) = E_old[ sum_t r_t(theta) (R-b) ] with r_t = pi_theta(o_t|s_t)/pi_old(o_t|s_t).\n",
+    "They agree at theta_old and drift apart as ||theta - theta_old|| grows: value error O(delta^2), gradient error O(delta)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "source": [
+    "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n",
+    "R = default_reward(lm)\n",
+    "th_old = lm.theta0.clone()\n",
+    "q = lm.probs(th_old)\n",
+    "b = float(q @ R)\n",
+    "def exact_grad(th):\n",
+    "    t = th.clone().requires_grad_(True)\n",
+    "    return torch.autograd.grad((lm.probs(t) * R).sum(), t)[0]\n",
+    "def seq_is_grad(th):\n",
+    "    t = th.clone().requires_grad_(True)\n",
+    "    lpn = lm.seq_logprob(t)\n",
+    "    wts = (lpn - lm.seq_logprob(th_old)).exp().detach()\n",
+    "    return torch.autograd.grad((q * wts * (R - b) * lpn).sum(), t)[0]\n",
+    "def token_surrogate(th):\n",
+    "    t = th.clone().requires_grad_(True)\n",
+    "    r = (lm.token_logprobs(t) - lm.token_logprobs(th_old)).exp() * lm.mask       # r_t(theta), 0 where masked\n",
+    "    L = (q[:, None] * r * (R - b)[:, None]).sum()\n",
+    "    return L, torch.autograd.grad(L, t)[0]\n",
+    "deltas = [0.0, 0.02, 0.04, 0.08, 0.16, 0.32]\n",
+    "L0 = float((q * lm.lengths * (R - b)).sum())\n",
+    "gerr, verr = [], []\n",
+    "for d in deltas:\n",
+    "    th = perturb(th_old, d, seed=5) if d > 0 else th_old.clone()\n",
+    "    ge = exact_grad(th)\n",
+    "    check(f\"sequence-level IS gradient == exact grad J, delta={d}\", maxabs(seq_is_grad(th), ge), 1e-14) if d in (0.0, 0.32) else None\n",
+    "    Lv, gt = token_surrogate(th)\n",
+    "    gerr.append(float((gt - ge).norm())); verr.append(abs(float(Lv.detach()) - L0 - float((lm.probs(th) * R).sum() - (q * R).sum())))\n",
+    "    print(f\"delta={d:5.2f}   |grad_tok - grad_J| = {gerr[-1]:.3e}   |dL - dJ| = {verr[-1]:.3e}\")\n",
+    "check(\"token surrogate gradient == exact gradient at theta_old\", gerr[0], 1e-14)\n",
+    "check(\"token surrogate value == J to first order at theta_old (error at delta=0)\", verr[0], 1e-14)\n",
+    "import numpy as np\n",
+    "sl_g = np.polyfit(np.log(deltas[1:]), np.log(gerr[1:]), 1)[0]\n",
+    "sl_v = np.polyfit(np.log(deltas[1:]), np.log(verr[1:]), 1)[0]\n",
+    "print(f\"log-log slopes: gradient error {sl_g:.2f} (expect ~1), value error {sl_v:.2f} (expect ~2)\")\n",
+    "check(\"gradient-error slope ~ 1\", abs(sl_g - 1.0), 0.15)\n",
+    "check(\"value-error slope ~ 2\", abs(sl_v - 2.0), 0.15)\n",
+    "print(\"done; checks:\", len(CHECKS))"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|sequence-level IS gradient == exact grad J, delta=0.0|4.163e-17|tol=1.0e-14|PASS\n",
+      "delta= 0.00   |grad_tok - grad_J| = 6.150e-17   |dL - dJ| = 2.776e-17\n",
+      "delta= 0.02   |grad_tok - grad_J| = 2.905e-04   |dL - dJ| = 3.869e-07\n",
+      "delta= 0.04   |grad_tok - grad_J| = 5.814e-04   |dL - dJ| = 1.550e-06\n",
+      "delta= 0.08   |grad_tok - grad_J| = 1.164e-03   |dL - dJ| = 6.214e-06\n",
+      "delta= 0.16   |grad_tok - grad_J| = 2.333e-03   |dL - dJ| = 2.497e-05\n",
+      "CHECK|exact|sequence-level IS gradient == exact grad J, delta=0.32|2.776e-17|tol=1.0e-14|PASS\n",
+      "delta= 0.32   |grad_tok - grad_J| = 4.684e-03   |dL - dJ| = 1.007e-04\n",
+      "CHECK|exact|token surrogate gradient == exact gradient at theta_old|6.150e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|token surrogate value == J to first order at theta_old (error at delta=0)|2.776e-17|tol=1.0e-14|PASS\n",
+      "log-log slopes: gradient error 1.00 (expect ~1), value error 2.01 (expect ~2)\n",
+      "CHECK|exact|gradient-error slope ~ 1|2.652e-03|tol=1.5e-01|PASS\n",
+      "CHECK|exact|value-error slope ~ 2|5.861e-03|tol=1.5e-01|PASS\n",
+      "done; checks: 23\n"
+     ]
+    }
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

code/notebooks/0104-trust-region-ppo-clip.ipynb

@@ -0,0 +1,318 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 0104 - Trust regions -> PPO clipped surrogate; the clip as a pessimistic bound; the gradient case table\n",
+    "Companion to `notes/0104 - Trust Regions and the PPO Clip.pdf`.  Exact enumeration, fp64."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "source": [
+    "import sys, math, itertools\n",
+    "sys.path.insert(0, \"..\")\n",
+    "from _m1_lib import *\n",
+    "import numpy as np\n",
+    "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n",
+    "R = default_reward(lm)\n",
+    "th_old = lm.theta0.clone()\n",
+    "q = lm.probs(th_old)\n",
+    "seqset = {s: i for i, s in enumerate(lm.seqs)}\n",
+    "V, L = lm.V, lm.L"
+   ],
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Exact value functions by backward recursion; advantage A_old(s,a) = V(s+a) - V(s)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "source": [
+    "from functools import lru_cache\n",
+    "def make_value(theta):\n",
+    "    @lru_cache(None)\n",
+    "    def Vf(prefix):\n",
+    "        if prefix in seqset: return float(R[seqset[prefix]])        # terminal (EOS or length L): value = reward\n",
+    "        cp = lm.cond_probs(theta, prefix)\n",
+    "        return float(sum(cp[a] * Vf(prefix + (a,)) for a in range(V)))\n",
+    "    return Vf\n",
+    "Vold = make_value(th_old)\n",
+    "check(\"V_old(root) = J(theta_old) = sum_o pi R\", abs(Vold(()) - float(q @ R)), 1e-14)\n",
+    "states = lm.prefixes()\n",
+    "def A_old(prefix, a): return Vold(prefix + (a,)) - Vold(prefix)\n",
+    "check(\"E_{a~pi_old}[A_old(s,a)] = 0 at every state\", max(abs(sum(float(lm.cond_probs(th_old, s)[a]) * A_old(s, a) for a in range(V))) for s in states), 1e-14)\n",
+    "# advantage along every enumerated sequence: Adv[o,t] = A_old(s_t, o_t)\n",
+    "Adv = torch.zeros(lm.N, L)\n",
+    "for i, s in enumerate(lm.seqs):\n",
+    "    for t in range(len(s)): Adv[i, t] = A_old(s[:t], s[t])\n",
+    "check(\"telescoping:  sum_t A_old(s_t,o_t) = R(o) - J(theta_old) for EVERY sequence\", maxabs(Adv.sum(1), R - Vold(())), 1e-14)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|V_old(root) = J(theta_old) = sum_o pi R|0.000e+00|tol=1.0e-14|PASS\n",
+      "CHECK|exact|E_{a~pi_old}[A_old(s,a)] = 0 at every state|1.475e-16|tol=1.0e-14|PASS\n",
+      "CHECK|exact|telescoping:  sum_t A_old(s_t,o_t) = R(o) - J(theta_old) for EVERY sequence|1.110e-16|tol=1.0e-14|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Performance-difference lemma:  J(theta') - J(theta_old) = E_{o~pi'}[ sum_t A_old(s_t,o_t) ]\n",
+    "and the exact IS form with the PREFIX ratio  w_{1:t} = prod_{s<=t} r_s   (action at t is the only thing that matters for A_old)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "source": [
+    "def ratios(th):\n",
+    "    lp, lq = lm.token_logprobs(th), lm.token_logprobs(th_old)\n",
+    "    r = (lp - lq).exp() * lm.mask + (1 - lm.mask)         # r_t (1 after EOS)\n",
+    "    return r, lp\n",
+    "for d in (0.1, 0.4):\n",
+    "    th = perturb(th_old, d, seed=3)\n",
+    "    p = lm.probs(th)\n",
+    "    dJ = float(p @ R - q @ R)\n",
+    "    pd_lemma = float((p * Adv.sum(1)).sum())                      # E_{pi'}[sum_t A_old]\n",
+    "    r, _ = ratios(th)\n",
+    "    wprefix = r.cumprod(1)                                         # w_{1:t}\n",
+    "    is_exact = float((q[:, None] * wprefix * Adv).sum())           # E_old[ sum_t w_{1:t} A_old(s_t,o_t) ]\n",
+    "    L_sur = float((q[:, None] * r * Adv).sum())                    # TRPO/PPO surrogate: only the token ratio\n",
+    "    check(f\"performance difference lemma, delta={d}\", abs(dJ - pd_lemma), 1e-14)\n",
+    "    check(f\"IS with prefix ratio w_(1:t) is exact, delta={d}\", abs(dJ - is_exact), 1e-14)\n",
+    "    print(f\"   delta={d}:  dJ={dJ:+.6f}   surrogate L-J_old={L_sur:+.6f}   (error {abs(dJ-L_sur):.2e})\")\n",
+    "# first-order agreement of surrogate and J at theta_old\n",
+    "th = th_old.clone().requires_grad_(True)\n",
+    "lpn, lqo = lm.token_logprobs(th), lm.token_logprobs(th_old)\n",
+    "Ls = (q[:, None] * ((lpn - lqo).exp() * lm.mask) * Adv).sum()\n",
+    "gL = torch.autograd.grad(Ls, th)[0]\n",
+    "th2 = th_old.clone().requires_grad_(True)\n",
+    "gJ = torch.autograd.grad((lm.probs(th2) * R).sum(), th2)[0]\n",
+    "check(\"grad of surrogate at theta_old == grad J (first-order matching)\", maxabs(gL, gJ), 1e-14)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|performance difference lemma, delta=0.1|4.857e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|IS with prefix ratio w_(1:t) is exact, delta=0.1|4.380e-17|tol=1.0e-14|PASS\n",
+      "   delta=0.1:  dJ=+0.003183   surrogate L-J_old=+0.003188   (error 5.85e-06)\n",
+      "CHECK|exact|performance difference lemma, delta=0.4|2.776e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|IS with prefix ratio w_(1:t) is exact, delta=0.4|6.939e-18|tol=1.0e-14|PASS\n",
+      "   delta=0.4:  dJ=+0.013143   surrogate L-J_old=+0.013234   (error 9.14e-05)\n",
+      "CHECK|exact|grad of surrogate at theta_old == grad J (first-order matching)|2.082e-17|tol=1.0e-14|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Finite-horizon trust-region bound (derived in the note):  |J' - L| <= 2 eps alpha^2 T (T-1)\n",
+    "alpha = max_s TV(pi'(.|s), pi(.|s)),  eps = max_{s,a} |A_old(s,a)|,  T = horizon.  Pinsker: alpha^2 <= D_KL^max / 2."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "source": [
+    "eps = max(abs(A_old(s, a)) for s in states for a in range(V))\n",
+    "print(\"eps = max|A| =\", eps)\n",
+    "rows = []\n",
+    "for d in (0.05, 0.1, 0.2, 0.4, 0.8, 1.6):\n",
+    "    th = perturb(th_old, d, seed=3)\n",
+    "    alpha = max(0.5 * float((lm.cond_probs(th, s) - lm.cond_probs(th_old, s)).abs().sum()) for s in states)\n",
+    "    kl_max = max(float((lm.cond_probs(th_old, s) * (lm.cond_probs(th_old, s).log() - lm.cond_probs(th, s).log())).sum()) for s in states)\n",
+    "    p = lm.probs(th); r, _ = ratios(th)\n",
+    "    err = abs(float(p @ R - q @ R) - float((q[:, None] * r * Adv).sum()))\n",
+    "    bound_tv = 2 * eps * alpha ** 2 * L * (L - 1)\n",
+    "    bound_kl = eps * L * (L - 1) * kl_max\n",
+    "    rows.append((d, alpha, kl_max, err, bound_tv, bound_kl))\n",
+    "    print(f\"delta={d:4.2f}  alpha={alpha:.4f}  KLmax={kl_max:.4f}  |J'-L|={err:.3e}  TV-bound={bound_tv:.3e}  KL-bound={bound_kl:.3e}  slack={bound_tv/max(err,1e-300):.1f}x\")\n",
+    "    assert err <= bound_tv + 1e-15 and bound_tv <= bound_kl + 1e-15 and alpha ** 2 <= kl_max / 2 + 1e-15\n",
+    "check(\"bound |J'-L| <= 2 eps a^2 T(T-1) holds at all 6 trust-region sizes (max err/bound)\", max(r[3] / r[4] for r in rows), 1.0)\n",
+    "check(\"Pinsker chain a^2 <= KLmax/2 holds (max a^2/(KLmax/2))\", max(r[1] ** 2 / (r[2] / 2) for r in rows), 1.0)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "eps = max|A| = 0.9030066306956115\n",
+      "delta=0.05  alpha=0.0048  KLmax=0.0001  |J'-L|=1.471e-06  TV-bound=5.034e-04  KL-bound=6.045e-04  slack=342.3x\n",
+      "delta=0.10  alpha=0.0097  KLmax=0.0002  |J'-L|=5.854e-06  TV-bound=2.028e-03  KL-bound=2.425e-03  slack=346.5x\n",
+      "delta=0.20  alpha=0.0195  KLmax=0.0009  |J'-L|=2.320e-05  TV-bound=8.226e-03  KL-bound=9.755e-03  slack=354.6x\n",
+      "delta=0.40  alpha=0.0395  KLmax=0.0036  |J'-L|=9.139e-05  TV-bound=3.379e-02  KL-bound=3.945e-02  slack=369.7x\n",
+      "delta=0.80  alpha=0.0808  KLmax=0.0149  |J'-L|=3.606e-04  TV-bound=1.415e-01  KL-bound=1.611e-01  slack=392.4x\n",
+      "delta=1.60  alpha=0.1668  KLmax=0.0615  |J'-L|=1.532e-03  TV-bound=6.027e-01  KL-bound=6.664e-01  slack=393.4x\n",
+      "CHECK|exact|bound |J'-L| <= 2 eps a^2 T(T-1) holds at all 6 trust-region sizes (max err/bound)|2.922e-03|tol=1.0e+00|PASS\n",
+      "CHECK|exact|Pinsker chain a^2 <= KLmax/2 holds (max a^2/(KLmax/2))|9.043e-01|tol=1.0e+00|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. PPO clip:  f(r, A) = min(r A, clip(r, 1-eps_l, 1+eps_h) A)\n",
+    "closed form  A>0: A*min(r, 1+eps_h);   A<0: A*max(r, 1-eps_l);   pessimistic: f <= r A;   gradient case table."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "source": [
+    "def f_naive(r, A, el, eh):\n",
+    "    clipped = min(max(r, 1 - el), 1 + eh)\n",
+    "    return min(r * A, clipped * A)\n",
+    "def f_closed(r, A, el, eh):\n",
+    "    return A * min(r, 1 + eh) if A > 0 else A * max(r, 1 - el)\n",
+    "el, eh = 0.2, 0.28\n",
+    "rg = np.concatenate([np.linspace(0.0, 3.0, 301), [0.8 + 1e-9, 1.28 - 1e-9]])\n",
+    "worst_cf, worst_pess = 0.0, -1e9\n",
+    "for A in (-2.0, -0.3, 0.3, 2.0):\n",
+    "    for r in rg:\n",
+    "        worst_cf = max(worst_cf, abs(f_naive(r, A, el, eh) - f_closed(r, A, el, eh)))\n",
+    "        worst_pess = max(worst_pess, f_naive(r, A, el, eh) - r * A)\n",
+    "check(\"closed forms A*min(r,1+eh) / A*max(r,1-el) == min/clip definition\", worst_cf, 1e-14)\n",
+    "check(\"pessimism: f(r,A) <= r*A everywhere (max of f - rA, must be <= 0)\", max(0.0, worst_pess), 0.0)\n",
+    "# autograd gradient wrt log pi  vs  hand case table\n",
+    "def grad_autograd(r, A, el, eh):\n",
+    "    lp = torch.tensor(math.log(r), requires_grad=True)          # log pi_theta ; old log prob = 0 so r = exp(lp)\n",
+    "    rr = lp.exp()\n",
+    "    f = torch.minimum(rr * A, torch.clamp(rr, 1 - el, 1 + eh) * A)\n",
+    "    return float(torch.autograd.grad(f, lp)[0])\n",
+    "def grad_table(r, A, el, eh):\n",
+    "    if A > 0: passes = r < 1 + eh\n",
+    "    else:     passes = r > 1 - el\n",
+    "    return A * r if passes else 0.0                              # d f / d log pi = A r  when the token is live\n",
+    "worst, n_zero, n_pass = 0.0, 0, 0\n",
+    "for A in (-2.0, -0.3, 0.3, 2.0):\n",
+    "    for r in [x for x in rg if x > 0 and abs(x - (1 - el)) > 1e-6 and abs(x - (1 + eh)) > 1e-6]:\n",
+    "        ga, gt = grad_autograd(r, A, el, eh), grad_table(r, A, el, eh)\n",
+    "        worst = max(worst, abs(ga - gt)); n_zero += gt == 0.0; n_pass += gt != 0.0\n",
+    "check(f\"gradient case table == autograd on {n_zero+n_pass} (r,A) points ({n_zero} silenced, {n_pass} live)\", worst, 1e-14)\n",
+    "# the 2x3 table printed explicitly\n",
+    "print(\"\\n  gradient wrt log pi of one token (eps_l=0.2, eps_h=0.28), '.' = zero\")\n",
+    "print(\"  region          A>0         A<0\")\n",
+    "for name, r in ((\"r < 1-eps_l\", 0.5), (\"in [1-eps_l,1+eps_h]\", 1.0), (\"r > 1+eps_h\", 2.0)):\n",
+    "    print(f\"  {name:<20s}\", \"  \".join(f\"{(grad_autograd(r, A, el, eh)/(A*r)):4.0f}*A*r\" if abs(grad_autograd(r, A, el, eh))>0 else \"   .     \" for A in (0.7, -0.7)))\n",
+    "# boundary tie behaviour in torch (documented, not relied upon)\n",
+    "print(\"  boundary tie  r=1+eps_h, A>0:\", grad_autograd(1 + eh, 1.0, el, eh), \" r=1-eps_l, A<0:\", grad_autograd(1 - el, -1.0, el, eh))\n",
+    "# naive vs vectorised implementation\n",
+    "gen = torch.Generator().manual_seed(1)\n",
+    "rr, AA = torch.rand(10000, generator=gen) * 3, torch.randn(10000, generator=gen)\n",
+    "fn = torch.tensor([f_naive(float(a), float(b), el, eh) for a, b in zip(rr, AA)])\n",
+    "fv = torch.where(AA > 0, AA * torch.clamp(rr, max=1 + eh), AA * torch.clamp(rr, min=1 - el))\n",
+    "check(\"naive scalar loop == vectorised torch.where implementation\", maxabs(fn, fv), 1e-14)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|closed forms A*min(r,1+eh) / A*max(r,1-el) == min/clip definition|0.000e+00|tol=1.0e-14|PASS\n",
+      "CHECK|exact|pessimism: f(r,A) <= r*A everywhere (max of f - rA, must be <= 0)|0.000e+00|tol=0.0e+00|PASS\n",
+      "CHECK|exact|gradient case table == autograd on 1192 (r,A) points (502 silenced, 690 live)|8.882e-16|tol=1.0e-14|PASS\n",
+      "\n",
+      "  gradient wrt log pi of one token (eps_l=0.2, eps_h=0.28), '.' = zero\n",
+      "  region          A>0         A<0\n",
+      "  r < 1-eps_l             1*A*r     .     \n",
+      "  in [1-eps_l,1+eps_h]    1*A*r     1*A*r\n",
+      "  r > 1+eps_h             .          1*A*r\n",
+      "  boundary tie  r=1+eps_h, A>0: 1.28  r=1-eps_l, A<0: -0.8\n",
+      "CHECK|exact|naive scalar loop == vectorised torch.where implementation|0.000e+00|tol=1.0e-14|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. The clip is a heuristic, not a trust region: ratios leave [1-eps,1+eps] during multi-step ascent\n",
+    "Full-batch (exact-expectation) ascent on E_old[ sum_t f(r_t, A_t) ] with A = R - E[R] (outcome advantage), 600 steps."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "source": [
+    "R_ = R; b = float(q @ R_)\n",
+    "A_out = (R_ - b)[:, None].expand(-1, L) * lm.mask\n",
+    "th = th_old.clone().requires_grad_(True)\n",
+    "opt = torch.optim.SGD([th], lr=0.3)\n",
+    "lq = lm.token_logprobs(th_old)\n",
+    "hist = []\n",
+    "for step in range(601):\n",
+    "    lp = lm.token_logprobs(th)\n",
+    "    r = ((lp - lq).exp() * lm.mask) + (1 - lm.mask)\n",
+    "    f = torch.where(A_out > 0, A_out * torch.clamp(r, max=1 + 0.2), A_out * torch.clamp(r, min=1 - 0.2))\n",
+    "    obj = (q[:, None] * f * lm.mask).sum()\n",
+    "    opt.zero_grad(); (-obj).backward(); opt.step()\n",
+    "    if step in (0, 10, 50, 100, 200, 400, 600):\n",
+    "        out_of = float((q[:, None] * (((r - 1).abs() > 0.2) & (lm.mask > 0))).sum() / (q * lm.lengths).sum())\n",
+    "        hist.append((step, float((r - 1).abs().detach().max()), out_of, float(obj.detach())))\n",
+    "print(\"step  max|r-1|   prob-weighted token fraction outside [0.8,1.2]   objective\")\n",
+    "for h in hist: print(f\"{h[0]:4d}  {h[1]:8.3f}   {h[2]:8.3f}   {h[3]:+.4f}\")\n",
+    "assert hist[-1][1] > 0.4\n",
+    "print(\"final max|r-1| =\", hist[-1][1], \"(eps=0.2)\")\n",
+    "print(\"done; checks:\", len(CHECKS))"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "step  max|r-1|   prob-weighted token fraction outside [0.8,1.2]   objective\n",
+      "   0     0.000      0.000   +0.2230\n",
+      "  10     0.252      0.113   +0.2838\n",
+      "  50     0.452      0.475   +0.2965\n",
+      " 100     0.558      0.530   +0.2998\n",
+      " 200     0.654      0.580   +0.3016\n",
+      " 400     0.751      0.753   +0.3021\n",
+      " 600     0.812      0.748   +0.3031\n",
+      "final max|r-1| = 0.8121891873787255 (eps=0.2)\n",
+      "done; checks: 14\n"
+     ]
+    }
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

code/notebooks/0105-grpo-group-advantages.ipynb

@@ -0,0 +1,403 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 0105 - GRPO: group-relative advantages as a Monte-Carlo baseline; std and length normalisation biases; Dr.GRPO\n",
+    "Companion to `notes/0105 - GRPO and Dr GRPO.pdf`.  Exact enumeration of ALL G-tuples of samples (fp64)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "source": [
+    "import sys, math\n",
+    "sys.path.insert(0, \"..\")\n",
+    "from _m1_lib import *\n",
+    "import numpy as np\n",
+    "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n",
+    "theta = lm.theta0.clone()\n",
+    "R = default_reward(lm)\n",
+    "p = lm.probs(theta)\n",
+    "S = score_matrix(lm, theta)                      # [N,P]\n",
+    "N = lm.N\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "pr = (lm.probs(th) * R).sum()                    # success probability p(theta)\n",
+    "grad_p = torch.autograd.grad(pr, th)[0]\n",
+    "p_succ = float(pr.detach())\n",
+    "print(\"success prob p =\", p_succ, \"  ||grad p|| =\", float(grad_p.norm()))"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "success prob p = 0.16592597945219614   ||grad p|| = 0.18246238486098645\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Advantage definitions, RLOO equivalence, invariances"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "source": [
+    "def adv_grpo(Rg, ddof=0, eps=0.0):\n",
+    "    mu = Rg.mean(); sd = Rg.std(unbiased=bool(ddof)) if len(Rg) > 1 else torch.tensor(0.0)\n",
+    "    return (Rg - mu) / (sd + eps)\n",
+    "def adv_center(Rg): return Rg - Rg.mean()\n",
+    "def adv_rloo(Rg):\n",
+    "    G = len(Rg); return Rg - (Rg.sum() - Rg) / (G - 1)\n",
+    "gen = torch.Generator().manual_seed(0)\n",
+    "Rg = torch.rand(7, generator=gen)\n",
+    "G_ = len(Rg)\n",
+    "check(\"G/(G-1) * (R_i - mean) == RLOO advantage\", maxabs(G_ / (G_ - 1) * adv_center(Rg), adv_rloo(Rg)), 1e-15)\n",
+    "check(\"GRPO advantage invariant to R -> a R + c (a>0)\", maxabs(adv_grpo(Rg), adv_grpo(3.7 * Rg - 2.0)), 1e-13)\n",
+    "check(\"centred advantage scales with a:  A(aR+c) = a A(R)\", maxabs(adv_center(3.7 * Rg - 2.0), 3.7 * adv_center(Rg)), 1e-14)\n",
+    "check(\"GRPO advantage has mean 0 and (population) std 1\", max(abs(float(adv_grpo(Rg).mean())), abs(float(adv_grpo(Rg).std(unbiased=False)) - 1)), 1e-14)\n",
+    "Rdead = torch.ones(5)\n",
+    "check(\"dead group (all rewards equal) has zero advantage (eps=1e-4 convention)\", float(adv_grpo(Rdead, eps=1e-4).abs().max()), 0.0)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|G/(G-1) * (R_i - mean) == RLOO advantage|1.110e-16|tol=1.0e-15|PASS\n",
+      "CHECK|exact|GRPO advantage invariant to R -> a R + c (a>0)|3.331e-16|tol=1.0e-13|PASS\n",
+      "CHECK|exact|centred advantage scales with a:  A(aR+c) = a A(R)|3.331e-16|tol=1.0e-14|PASS\n",
+      "CHECK|exact|GRPO advantage has mean 0 and (population) std 1|1.269e-16|tol=1.0e-14|PASS\n",
+      "CHECK|exact|dead group (all rewards equal) has zero advantage (eps=1e-4 convention)|0.000e+00|tol=0.0e+00|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. EXACT expected gradient of the group estimators, by enumerating every G-tuple of responses\n",
+    "  estimator:  g = (1/G) sum_i A_i * score(o_i)       (outcome reward, one scalar advantage per response)\n",
+    "compare:  mean-centred (Dr.GRPO)  ->  (1 - 1/G) grad J ;   std-normalised (GRPO) -> c_G(p) grad p  with c_G NOT constant."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "source": [
+    "def exact_group_gradient(G, kind, Rr=None):\n",
+    "    \"\"\"Sum over all N^G tuples with product probability.  kind in {'center','grpo','rloo'}.\"\"\"\n",
+    "    Rr = R if Rr is None else Rr\n",
+    "    assert G in (2, 3)\n",
+    "    if G == 2:\n",
+    "        R1, R2 = Rr[:, None].expand(N, N), Rr[None, :].expand(N, N)\n",
+    "        P = p[:, None] * p[None, :]\n",
+    "        A = _adv_batch(torch.stack([R1, R2]), kind)                      # [2, N, N]\n",
+    "        g = (P * A[0]).sum(1) @ S + (P * A[1]).sum(0) @ S                # sample 1 carries score(o1), sample 2 score(o2)\n",
+    "        return g / 2\n",
+    "    else:\n",
+    "        g = torch.zeros(S.shape[1])\n",
+    "        R2, R3 = Rr[:, None], Rr[None, :]\n",
+    "        P23 = p[:, None] * p[None, :]\n",
+    "        for i in range(N):\n",
+    "            Rs = torch.stack([Rr[i].expand(N, N), R2.expand(N, N), R3.expand(N, N)])\n",
+    "            A = _adv_batch(Rs, kind)                                     # [3, N, N]\n",
+    "            g += p[i] * (P23 * A[0]).sum() * S[i]\n",
+    "            g += p[i] * ((P23 * A[1]).sum(1) @ S)\n",
+    "            g += p[i] * ((P23 * A[2]).sum(0) @ S)\n",
+    "        return g / 3\n",
+    "def _adv_batch(Rs, kind):\n",
+    "    \"\"\"Rs [G, ...] rewards of the group, vectorised over the trailing dims.\"\"\"\n",
+    "    G = Rs.shape[0]\n",
+    "    mu = Rs.mean(0, keepdim=True)\n",
+    "    if kind == 'center': return Rs - mu\n",
+    "    if kind == 'rloo': return Rs - (Rs.sum(0, keepdim=True) - Rs) / (G - 1)\n",
+    "    sd = ((Rs - mu) ** 2).mean(0, keepdim=True).sqrt()                  # population std\n",
+    "    safe = torch.where(sd > 0, sd, torch.ones_like(sd))\n",
+    "    return torch.where(sd > 0, (Rs - mu) / safe, torch.zeros_like(Rs))\n",
+    "for G in (2, 3):\n",
+    "    gc = exact_group_gradient(G, 'center')\n",
+    "    gr = exact_group_gradient(G, 'rloo')\n",
+    "    check(f\"G={G}: mean-centred estimator == (1-1/G) grad J exactly (Dr.GRPO unbiasedness)\", maxabs(gc, (1 - 1 / G) * grad_p), 1e-13)\n",
+    "    check(f\"G={G}: RLOO estimator == grad J exactly\", maxabs(gr, grad_p), 1e-13)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|G=2: mean-centred estimator == (1-1/G) grad J exactly (Dr.GRPO unbiasedness)|1.388e-17|tol=1.0e-13|PASS\n",
+      "CHECK|exact|G=2: RLOO estimator == grad J exactly|2.776e-17|tol=1.0e-13|PASS\n",
+      "CHECK|exact|G=3: mean-centred estimator == (1-1/G) grad J exactly (Dr.GRPO unbiasedness)|2.776e-17|tol=1.0e-13|PASS\n",
+      "CHECK|exact|G=3: RLOO estimator == grad J exactly|4.857e-17|tol=1.0e-13|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Std normalisation: for binary rewards the expected gradient is c_G(p) * grad p with an explicit c_G\n",
+    "k successes among G (0<k<G): A_succ = sqrt((G-k)/k), A_fail = -sqrt(k/(G-k)); dead groups give 0.\n",
+    "By exchangeability E[g] = (a1 - a0) grad p  with a1 = E[A_1 | R_1=1], a0 = E[A_1 | R_1=0]."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "source": [
+    "def binom_pmf(n, k, pp): return math.comb(n, k) * pp ** k * (1 - pp) ** (n - k)\n",
+    "def c_G(G, pp):\n",
+    "    a1 = sum(binom_pmf(G - 1, K, pp) * (math.sqrt((G - (1 + K)) / (1 + K)) if 1 + K < G else 0.0) for K in range(G))\n",
+    "    a0 = sum(binom_pmf(G - 1, K, pp) * (-math.sqrt(K / (G - K)) if 1 <= K < G else 0.0) for K in range(G))\n",
+    "    return a1 - a0\n",
+    "for G in (2, 3):\n",
+    "    gg = exact_group_gradient(G, 'grpo')\n",
+    "    check(f\"G={G}: brute-force GRPO(std-norm) gradient == c_G(p) * grad p\", maxabs(gg, c_G(G, p_succ) * grad_p), 1e-13)\n",
+    "# G = 4 by exchangeability over the 3 other samples (still an enumeration of N^3 = 1.77M tuples)\n",
+    "G = 4; Rr = R\n",
+    "R2, R3, R4 = Rr[:, None, None], Rr[None, :, None], Rr[None, None, :]\n",
+    "P234 = p[:, None, None] * p[None, :, None] * p[None, None, :]\n",
+    "Wr = {}\n",
+    "for r1 in (0.0, 1.0):\n",
+    "    Rs = torch.stack([torch.full((N, N, N), r1), R2.expand(N, N, N), R3.expand(N, N, N), R4.expand(N, N, N)])\n",
+    "    Wr[r1] = float((P234 * _adv_batch(Rs, 'grpo')[0]).sum())\n",
+    "g4 = (Wr[1.0] * (p * R) @ S) + (Wr[0.0] * (p * (1 - R)) @ S)\n",
+    "check(\"G=4: enumerated (N^3 other samples) GRPO gradient == c_4(p) grad p\", maxabs(g4, c_G(4, p_succ) * grad_p), 1e-13)\n",
+    "print(\"\\n  c_G(p) at p = %.3f for various G, vs 1/sqrt(p(1-p)) = %.4f and Dr.GRPO's (1-1/G):\" % (p_succ, 1 / math.sqrt(p_succ * (1 - p_succ))))\n",
+    "for G in (2, 4, 8, 16, 64, 256):\n",
+    "    print(f\"   G={G:4d}   c_G = {c_G(G, p_succ):.4f}\")\n",
+    "check(\"c_G(p) -> 1/sqrt(p(1-p)) as G grows (G=256, relative error)\", abs(c_G(256, p_succ) * math.sqrt(p_succ * (1 - p_succ)) - 1), 0.1)\n",
+    "check(\"c_G(p->0) = sqrt(G-1) (bounded up-weighting of hard prompts), G=16\", abs(c_G(16, 1e-9) - math.sqrt(15)), 1e-6)\n",
+    "print(\"   difficulty weighting c_G(p)  (G=16):\")\n",
+    "for pp in (0.01, 0.05, 0.2, 0.5, 0.8, 0.95, 0.99):\n",
+    "    print(f\"     p={pp:4.2f}   GRPO c = {c_G(16, pp):6.3f}    Dr.GRPO c = {1 - 1/16:.3f}\")\n",
+    "# the infinite-G limit is the gradient of the variance-stabilising objective 2*arcsin(sqrt(p))\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "pp_ = (lm.probs(th) * R).sum()\n",
+    "g_asin = torch.autograd.grad(2 * torch.asin(pp_.sqrt()), th)[0]\n",
+    "check(\"grad[2 asin(sqrt p)] == grad p / sqrt(p(1-p))\", maxabs(g_asin, grad_p / math.sqrt(p_succ * (1 - p_succ))), 1e-14)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|G=2: brute-force GRPO(std-norm) gradient == c_G(p) * grad p|2.776e-17|tol=1.0e-13|PASS\n",
+      "CHECK|exact|G=3: brute-force GRPO(std-norm) gradient == c_G(p) * grad p|5.551e-17|tol=1.0e-13|PASS\n",
+      "CHECK|exact|G=4: enumerated (N^3 other samples) GRPO gradient == c_4(p) grad p|8.327e-17|tol=1.0e-13|PASS\n",
+      "\n",
+      "  c_G(p) at p = 0.166 for various G, vs 1/sqrt(p(1-p)) = 2.6881 and Dr.GRPO's (1-1/G):\n",
+      "   G=   2   c_G = 1.0000\n",
+      "   G=   4   c_G = 1.6678\n",
+      "   G=   8   c_G = 2.1725\n",
+      "   G=  16   c_G = 2.4801\n",
+      "   G=  64   c_G = 2.6487\n",
+      "   G= 256   c_G = 2.6785\n",
+      "CHECK|exact|c_G(p) -> 1/sqrt(p(1-p)) as G grows (G=256, relative error)|3.559e-03|tol=1.0e-01|PASS\n",
+      "CHECK|exact|c_G(p->0) = sqrt(G-1) (bounded up-weighting of hard prompts), G=16|1.454e-08|tol=1.0e-06|PASS\n",
+      "   difficulty weighting c_G(p)  (G=16):\n",
+      "     p=0.01   GRPO c =  3.733    Dr.GRPO c = 0.938\n",
+      "     p=0.05   GRPO c =  3.267    Dr.GRPO c = 0.938\n",
+      "     p=0.20   GRPO c =  2.346    Dr.GRPO c = 0.938\n",
+      "     p=0.50   GRPO c =  1.934    Dr.GRPO c = 0.938\n",
+      "     p=0.80   GRPO c =  2.346    Dr.GRPO c = 0.938\n",
+      "     p=0.95   GRPO c =  3.267    Dr.GRPO c = 0.938\n",
+      "     p=0.99   GRPO c =  3.733    Dr.GRPO c = 0.938\n",
+      "CHECK|exact|grad[2 asin(sqrt p)] == grad p / sqrt(p(1-p))|2.776e-17|tol=1.0e-14|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3b. General principle: ANY group advantage  A_1 = h(R_1; R_2..R_G)  has expected gradient  E[ f(R) score ]  with  f(r) = E_others[ h(r; R_2..R_G) ]\n",
+    "i.e. it is exactly the policy gradient of a *reshaped reward* f(R).  Check with a graded (3-valued) reward, G=3, std-normalised."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "source": [
+    "Rgr = R + (lm.lengths == 4).double()                  # rewards in {0,1,2}\n",
+    "g_graded = exact_group_gradient(3, 'grpo', Rr=Rgr)\n",
+    "vals = sorted(set(Rgr.tolist()))\n",
+    "f_of = {}\n",
+    "for r1 in vals:                                        # f(r) = E over the other two samples of A_1(r; R_2, R_3)\n",
+    "    R2, R3 = Rgr[:, None].expand(N, N), Rgr[None, :].expand(N, N)\n",
+    "    P23 = p[:, None] * p[None, :]\n",
+    "    A1 = _adv_batch(torch.stack([torch.full((N, N), r1), R2, R3]), 'grpo')[0]\n",
+    "    f_of[r1] = float((P23 * A1).sum())\n",
+    "f_vec = torch.tensor([f_of[float(x)] for x in Rgr])\n",
+    "check(\"graded reward, G=3: E[GRPO gradient] == E[f(R) score] with f(r)=E[A_1|R_1=r]\", maxabs(g_graded, (p[:, None] * f_vec[:, None] * S).sum(0)), 1e-13)\n",
+    "print(\"   reshaped reward f(r):\", {k: round(v, 4) for k, v in f_of.items()}, \"(true rewards:\", vals, \")\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|graded reward, G=3: E[GRPO gradient] == E[f(R) score] with f(r)=E[A_1|R_1=r]|5.551e-17|tol=1.0e-13|PASS\n",
+      "   reshaped reward f(r): {0.0: -0.3371, 1.0: 0.9871, 2.0: 1.2765} (true rewards: [0.0, 1.0, 2.0] )\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Length normalisation:  E[(R-b)/|o| * score] = grad E[(R-b)/|o|]  (b held fixed) -- a different objective"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "source": [
+    "b = float(p @ R)\n",
+    "lens = lm.lengths\n",
+    "g_unb = (p[:, None] * (R - b)[:, None] * S).sum(0)                       # = grad J\n",
+    "g_len = (p[:, None] * ((R - b) / lens)[:, None] * S).sum(0)               # GRPO-style 1/|o|\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "Jlen = (lm.probs(th) * (R - b) / lens).sum()                              # E[(R-b)/|o|] with b constant\n",
+    "g_Jlen = torch.autograd.grad(Jlen, th)[0]\n",
+    "check(\"unbiased baseline estimator == grad J\", maxabs(g_unb, grad_p), 1e-14)\n",
+    "check(\"1/|o|-normalised estimator == grad of E[(R-b)/|o|] (b fixed)\", maxabs(g_len, g_Jlen), 1e-14)\n",
+    "check(\"constant normaliser 1/Lmax gives grad J / Lmax (Dr.GRPO)\", maxabs((p[:, None] * (R - b)[:, None] * S).sum(0) / lm.L * lm.L, grad_p), 1e-14)\n",
+    "# effective per-sequence weight and which way the update pushes the expected length\n",
+    "Elen = (lm.probs(th) * lens).sum()\n",
+    "g_len_E = torch.autograd.grad(Elen, th)[0]                                # grad E[|o|]\n",
+    "print(\"\\n  length |  P(|o|)   mean weight on FAILURES  (unbiased | 1/|o|)    mean weight on SUCCESSES (unbiased | 1/|o|)\")\n",
+    "for Lk in range(1, lm.L + 1):\n",
+    "    m = lens == Lk\n",
+    "    pm = float(p[m].sum())\n",
+    "    fm, sm = m & (R == 0), m & (R == 1)\n",
+    "    wf = (-b, -b / Lk) ; ws = (1 - b, (1 - b) / Lk)\n",
+    "    print(f\"   {Lk}    | {pm:6.3f}   {wf[0]:+.3f} | {wf[1]:+.3f}  (mass {float(p[fm].sum()):.3f})        {ws[0]:+.3f} | {ws[1]:+.3f}  (mass {float(p[sm].sum()):.3f})\")\n",
+    "def part(wfun, sel):\n",
+    "    return float(((p * sel * wfun)[:, None] * S).sum(0) @ g_len_E)\n",
+    "fail, succ = (R == 0).double(), (R == 1).double()\n",
+    "print(\"  first-order d E|o| contributions  <grad E|o|, part of gradient>:\")\n",
+    "print(f\"     failures : unbiased {part(-b*torch.ones(N), fail):+.5f}   1/|o|-normalised {part(-b/lens, fail):+.5f}\")\n",
+    "print(f\"     successes: unbiased {part((1-b)*torch.ones(N), succ):+.5f}   1/|o|-normalised {part((1-b)/lens, succ):+.5f}\")\n",
+    "dir_unb = float(g_unb @ g_len_E); dir_len = float(g_len @ g_len_E)\n",
+    "print(f\"\\n  <grad E|o|, g_unbiased> = {dir_unb:+.5f}    <grad E|o|, g_1/|o|> = {dir_len:+.5f}   (first-order change in expected length per unit step)\")\n",
+    "# bias term = -b * grad E[1/|o|] + (grad E[R/|o|] - ... ) : decomposition check\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "gR_over = torch.autograd.grad((lm.probs(th) * R / lens).sum(), th)[0]\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "g_inv = torch.autograd.grad((lm.probs(th) / lens).sum(), th)[0]\n",
+    "check(\"decomposition g_len = grad E[R/|o|] - b grad E[1/|o|]\", maxabs(g_len, gR_over - b * g_inv), 1e-14)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|unbiased baseline estimator == grad J|2.776e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|1/|o|-normalised estimator == grad of E[(R-b)/|o|] (b fixed)|6.939e-18|tol=1.0e-14|PASS\n",
+      "CHECK|exact|constant normaliser 1/Lmax gives grad J / Lmax (Dr.GRPO)|2.776e-17|tol=1.0e-14|PASS\n",
+      "\n",
+      "  length |  P(|o|)   mean weight on FAILURES  (unbiased | 1/|o|)    mean weight on SUCCESSES (unbiased | 1/|o|)\n",
+      "   1    |  0.464   -0.166 | -0.166  (mass 0.464)        +0.834 | +0.834  (mass 0.000)\n",
+      "   2    |  0.258   -0.166 | -0.083  (mass 0.246)        +0.834 | +0.417  (mass 0.012)\n",
+      "   3    |  0.142   -0.166 | -0.055  (mass 0.049)        +0.834 | +0.278  (mass 0.094)\n",
+      "   4    |  0.136   -0.166 | -0.041  (mass 0.075)        +0.834 | +0.209  (mass 0.061)\n",
+      "  first-order d E|o| contributions  <grad E|o|, part of gradient>:\n",
+      "     failures : unbiased +0.02092   1/|o|-normalised +0.04135\n",
+      "     successes: unbiased +0.10518   1/|o|-normalised +0.03104\n",
+      "\n",
+      "  <grad E|o|, g_unbiased> = +0.12611    <grad E|o|, g_1/|o|> = +0.07239   (first-order change in expected length per unit step)\n",
+      "CHECK|exact|decomposition g_len = grad E[R/|o|] - b grad E[1/|o|]|7.806e-18|tol=1.0e-14|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Naive vs efficient loss reductions on a sampled group (gradient agreement) -- GRPO, DAPO-token, Dr.GRPO"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "source": [
+    "Gs = 6\n",
+    "gen = torch.Generator().manual_seed(42)\n",
+    "idx = torch.multinomial(p, Gs, replacement=True, generator=gen)\n",
+    "Rg = R[idx]\n",
+    "Ahat = adv_grpo(Rg, eps=1e-4); Actr = adv_center(Rg)\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "lp_tok = lm.token_logprobs(th)[idx]                    # [G, L]\n",
+    "m = lm.mask[idx]\n",
+    "lens_g = m.sum(1)\n",
+    "def surrogate_naive(lp_tok, A, mode):\n",
+    "    tot = 0.0\n",
+    "    if mode == 'grpo':\n",
+    "        for i in range(Gs): tot = tot + sum(A[i] * lp_tok[i, t] for t in range(int(lens_g[i]))) / lens_g[i]\n",
+    "        return tot / Gs\n",
+    "    if mode == 'dapo':\n",
+    "        n = float(lens_g.sum())\n",
+    "        for i in range(Gs): tot = tot + sum(A[i] * lp_tok[i, t] for t in range(int(lens_g[i])))\n",
+    "        return tot / n\n",
+    "    if mode == 'dr':\n",
+    "        for i in range(Gs): tot = tot + sum(A[i] * lp_tok[i, t] for t in range(int(lens_g[i])))\n",
+    "        return tot / (Gs * lm.L)\n",
+    "def surrogate_vec(lp_tok, A, mode):\n",
+    "    w = A[:, None] * m\n",
+    "    if mode == 'grpo': return ((w * lp_tok).sum(1) / lens_g).mean()\n",
+    "    if mode == 'dapo': return (w * lp_tok).sum() / m.sum()\n",
+    "    if mode == 'dr': return (w * lp_tok).sum() / (Gs * lm.L)\n",
+    "for mode, A in (('grpo', Ahat), ('dapo', Ahat), ('dr', Actr)):\n",
+    "    gn = torch.autograd.grad(surrogate_naive(lp_tok, A, mode), th, retain_graph=True)[0]\n",
+    "    gv = torch.autograd.grad(surrogate_vec(lp_tok, A, mode), th, retain_graph=True)[0]\n",
+    "    check(f\"naive python loops == vectorised masked reduction ({mode})\", maxabs(gn, gv), 1e-14)\n",
+    "print(\"done; checks:\", len(CHECKS))"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|naive python loops == vectorised masked reduction (grpo)|0.000e+00|tol=1.0e-14|PASS\n",
+      "CHECK|exact|naive python loops == vectorised masked reduction (dapo)|0.000e+00|tol=1.0e-14|PASS\n",
+      "CHECK|exact|naive python loops == vectorised masked reduction (dr)|0.000e+00|tol=1.0e-14|PASS\n",
+      "done; checks: 23\n"
+     ]
+    }
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

code/notebooks/0106-dapo-pathologies.ipynb

@@ -0,0 +1,450 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 0106 - DAPO: each of its four techniques derived as the fix for a named pathology\n",
+    "Companion to `notes/0106 - DAPO.pdf`.  fp64; exact enumeration where possible, Monte Carlo (flagged) elsewhere."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "source": [
+    "import sys, math\n",
+    "sys.path.insert(0, \"..\")\n",
+    "from _m1_lib import *\n",
+    "import numpy as np"
+   ],
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Clip-Higher (pathology: entropy collapse / exploration bottleneck)\n",
+    "1a. Asymmetric clip gradient table (eps_lo = 0.2, eps_hi = 0.28 as in DAPO Sec. 4.1) vs autograd; 1b. absolute probability headroom."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "source": [
+    "el, eh = 0.2, 0.28\n",
+    "def f_closed(r, A, el, eh): return A * min(r, 1 + eh) if A > 0 else A * max(r, 1 - el)\n",
+    "def grad_auto(r, A):\n",
+    "    lp = torch.tensor(math.log(r), requires_grad=True); rr = lp.exp()\n",
+    "    f = torch.minimum(rr * A, torch.clamp(rr, 1 - el, 1 + eh) * A)\n",
+    "    return float(torch.autograd.grad(f, lp)[0])\n",
+    "worst = 0.0; n = 0\n",
+    "for A in (-1.5, -0.2, 0.2, 1.5):\n",
+    "    for r in np.linspace(0.05, 3.0, 400):\n",
+    "        if min(abs(r - (1 - el)), abs(r - (1 + eh))) < 1e-6: continue\n",
+    "        live = (r < 1 + eh) if A > 0 else (r > 1 - el)\n",
+    "        worst = max(worst, abs(grad_auto(r, A) - (A * r if live else 0.0))); n += 1\n",
+    "check(f\"asymmetric clip: autograd == table (A r on live set) on {n} points\", worst, 1e-14)\n",
+    "print(\"   log-ratio limits: log(1+eps_hi)=%.4f  log(1-eps_lo)=%.4f  (symmetric eps=.2: +%.4f / %.4f)\" % (math.log(1 + eh), math.log(1 - el), math.log(1.2), math.log(0.8)))\n",
+    "print(\"   absolute probability reachable by ONE clipped update of a token with pi_old = p (A > 0):\")\n",
+    "print(\"      p        cap eps=0.20   cap eps=0.28    headroom(0.20)  headroom(0.28)\")\n",
+    "for pp in (0.001, 0.01, 0.1, 0.5, 0.9):\n",
+    "    c1, c2 = min(1.0, 1.2 * pp), min(1.0, 1.28 * pp)\n",
+    "    print(f\"     {pp:5.3f}     {c1:8.4f}      {c2:8.4f}        {c1-pp:8.4f}       {c2-pp:8.4f}\")\n",
+    "check(\"DAPO example: 0.01 -> 0.012 at eps=0.2\", abs(min(1, 1.2 * 0.01) - 0.012), 1e-15)\n",
+    "check(\"DAPO example: 0.9 -> min(1, 1.08) = 1.0 (no effective cap)\", abs(min(1, 1.2 * 0.9) - 1.0), 1e-15)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|asymmetric clip: autograd == table (A r on live set) on 1600 points|8.882e-16|tol=1.0e-14|PASS\n",
+      "   log-ratio limits: log(1+eps_hi)=0.2469  log(1-eps_lo)=-0.2231  (symmetric eps=.2: +0.1823 / -0.2231)\n",
+      "   absolute probability reachable by ONE clipped update of a token with pi_old = p (A > 0):\n",
+      "      p        cap eps=0.20   cap eps=0.28    headroom(0.20)  headroom(0.28)\n",
+      "     0.001       0.0012        0.0013          0.0002         0.0003\n",
+      "     0.010       0.0120        0.0128          0.0020         0.0028\n",
+      "     0.100       0.1200        0.1280          0.0200         0.0280\n",
+      "     0.500       0.6000        0.6400          0.1000         0.1400\n",
+      "     0.900       1.0000        1.0000          0.1000         0.1000\n",
+      "CHECK|exact|DAPO example: 0.01 -> 0.012 at eps=0.2|0.000e+00|tol=1.0e-15|PASS\n",
+      "CHECK|exact|DAPO example: 0.9 -> min(1, 1.08) = 1.0 (no effective cap)|0.000e+00|tol=1.0e-15|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "1c. Entropy dynamics of a softmax policy:  dH = -Cov_{a~pi}( log pi_a , dz_a )  (exact first order).  PG: dz = eta*pi*A;  NPG: dz = eta*A."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "source": [
+    "def entropy(z):\n",
+    "    lp = torch.log_softmax(z, 0); return float(-(lp.exp() * lp).sum())\n",
+    "def cov_formula(z, dz):\n",
+    "    pi = torch.softmax(z, 0); lg = pi.log()\n",
+    "    return float((pi * lg * dz).sum() - (pi * lg).sum() * (pi * dz).sum())\n",
+    "gen = torch.Generator().manual_seed(3)\n",
+    "worst_rel = 0.0\n",
+    "for trial in range(6):\n",
+    "    K = 5\n",
+    "    z = torch.randn(K, generator=gen) * 1.5\n",
+    "    A = torch.randn(K, generator=gen)\n",
+    "    pi = torch.softmax(z, 0); A = A - (pi * A).sum()                  # advantages must have zero pi-mean\n",
+    "    for name, dz in ((\"PG\", pi * A), (\"NPG\", A)):\n",
+    "        eta = 1e-6\n",
+    "        dH_fd = (entropy(z + eta * dz) - entropy(z - eta * dz)) / (2 * eta)\n",
+    "        dH_cov = -cov_formula(z, dz)\n",
+    "        worst_rel = max(worst_rel, abs(dH_fd - dH_cov) / max(1e-12, abs(dH_cov)))\n",
+    "check(\"entropy change = -Cov(log pi, dz): finite difference vs formula (PG & NPG, 6 random policies)\", worst_rel, 1e-6)\n",
+    "pi = torch.tensor([0.90, 0.05, 0.03, 0.02]); z = pi.log()\n",
+    "def dH_for(A):\n",
+    "    A = torch.tensor(A); A = A - (pi * A).sum(); return -cov_formula(z, pi * A)\n",
+    "print(\"   dH/d eta for tabular PG:  good action is the DOMINANT one (pi=.90): %+.5f   good action is a RARE one (pi=.02): %+.5f\" %\n",
+    "      (dH_for([1., 0, 0, 0]), dH_for([0., 0, 0, 1.])))\n",
+    "assert dH_for([1., 0, 0, 0]) < 0 < dH_for([0., 0, 0, 1.])\n",
+    "check(\"positive-advantage dominant token lowers entropy; positive-advantage rare token raises it (signs)\", 0.0, 0.0)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|entropy change = -Cov(log pi, dz): finite difference vs formula (PG & NPG, 6 random policies)|6.651e-09|tol=1.0e-06|PASS\n",
+      "   dH/d eta for tabular PG:  good action is the DOMINANT one (pi=.90): -0.03566   good action is a RARE one (pi=.02): +0.00641\n",
+      "CHECK|exact|positive-advantage dominant token lowers entropy; positive-advantage rare token raises it (signs)|0.000e+00|tol=0.0e+00|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "1d. Exact multi-step dynamics on a 2-action bandit (rare correct action a, pi=0.01): how fast can p(a) grow per rollout iteration?\n",
+    "Objective per inner step = sum_a pi_old(a) * f(r_a, A_a)  (exact expectation), tabular logits, plain gradient ascent."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "source": [
+    "def run_bandit(p0, eps_hi, outer=60, inner=400, lr=0.4, el=0.2):\n",
+    "    z = torch.log(torch.tensor([p0, 1 - p0])).clone()\n",
+    "    traj = [p0]\n",
+    "    for it in range(outer):\n",
+    "        pi_old = torch.softmax(z, 0).detach()\n",
+    "        A = torch.tensor([1.0, 0.0]) - pi_old[0]               # A_a = 1 - p, A_b = -p  (binary reward on action a)\n",
+    "        zz = z.clone().requires_grad_(True)\n",
+    "        for _ in range(inner):\n",
+    "            pi = torch.softmax(zz, 0); r = pi / pi_old\n",
+    "            f = torch.where(A > 0, A * torch.clamp(r, max=1 + eps_hi), A * torch.clamp(r, min=1 - el))\n",
+    "            obj = (pi_old * f).sum()\n",
+    "            g = torch.autograd.grad(obj, zz)[0]\n",
+    "            with torch.no_grad(): zz += lr * g\n",
+    "        z = zz.detach()\n",
+    "        traj.append(float(torch.softmax(z, 0)[0]))\n",
+    "    return torch.tensor(traj)\n",
+    "t20, t28, tinf = run_bandit(0.01, 0.2), run_bandit(0.01, 0.28), run_bandit(0.01, 10.0)\n",
+    "def first_cross(t, v=0.5):\n",
+    "    idx = (t >= v).nonzero()\n",
+    "    return int(idx[0]) if len(idx) else None\n",
+    "print(\"   iterations for p(a) to reach 0.5 starting from 0.01:  eps_hi=0.20 ->\", first_cross(t20), \"  eps_hi=0.28 ->\", first_cross(t28), \"  no upper clip ->\", first_cross(tinf))\n",
+    "g20 = (t20[1:6] / t20[:5]); g28 = (t28[1:6] / t28[:5])\n",
+    "print(\"   per-iteration growth factor p'/p over first 5 iterations: eps=.20:\", [round(float(x), 3) for x in g20], \" eps=.28:\", [round(float(x), 3) for x in g28])\n",
+    "print(\"   growth in iteration 1 slightly EXCEEDS 1+eps (1.251 vs 1.2; 1.342 vs 1.28): the clip stops only the DIRECT push on the rare action; the unclipped negative-advantage action keeps moving mass to it\")\n",
+    "assert first_cross(t28) <= first_cross(t20)\n",
+    "check(\"larger eps_hi never slows the rare action (iterations to 0.5: eps_hi=.28 <= eps_hi=.20)\", 0.0, 0.0)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "   iterations for p(a) to reach 0.5 starting from 0.01:  eps_hi=0.20 -> 8   eps_hi=0.28 -> 7   no upper clip -> 2\n",
+      "   per-iteration growth factor p'/p over first 5 iterations: eps=.20: [1.251, 1.283, 1.331, 1.454, 1.847]  eps=.28: [1.342, 1.389, 1.488, 1.805, 4.8]\n",
+      "   growth in iteration 1 slightly EXCEEDS 1+eps (1.251 vs 1.2; 1.342 vs 1.28): the clip stops only the DIRECT push on the rare action; the unclipped negative-advantage action keeps moving mass to it\n",
+      "CHECK|exact|larger eps_hi never slows the rare action (iterations to 0.5: eps_hi=.28 <= eps_hi=.20)|0.000e+00|tol=0.0e+00|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Dynamic sampling (pathology: zero-advantage groups).  P(dead group) = p^G + (1-p)^G"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "source": [
+    "torch.manual_seed(8)                                   # torch.distributions samplers use the global RNG\n",
+    "gen = torch.Generator().manual_seed(8)\n",
+    "for G in (4, 16):\n",
+    "    for pp in (0.02, 0.3):\n",
+    "        ks = torch.distributions.Binomial(G, torch.tensor(pp)).sample((400000,))\n",
+    "        emp = float(((ks == 0) | (ks == G)).double().mean()); ex = pp ** G + (1 - pp) ** G\n",
+    "        se = math.sqrt(ex * (1 - ex) / 400000)\n",
+    "        check(f\"P(dead) MC vs p^G+(1-p)^G  (G={G}, p={pp})  |diff|/SE\", abs(emp - ex) / se, 4.0, kind=\"mc\")\n",
+    "# expected live fraction for p ~ Beta(a,b):  1 - B(a+G,b)/B(a,b) - B(a,b+G)/B(a,b)\n",
+    "def lbeta(a, b): return math.lgamma(a) + math.lgamma(b) - math.lgamma(a + b)\n",
+    "def live_fraction(a, b, G): return 1 - math.exp(lbeta(a + G, b) - lbeta(a, b)) - math.exp(lbeta(a, b + G) - lbeta(a, b))\n",
+    "xs, ws = np.polynomial.legendre.leggauss(400); xs = 0.5 * (xs + 1); ws = 0.5 * ws\n",
+    "def live_quad(a, b, G):\n",
+    "    dens = np.exp(-lbeta(a, b)) * xs ** (a - 1) * (1 - xs) ** (b - 1)\n",
+    "    return float((ws * dens * (1 - xs ** G - (1 - xs) ** G)).sum())\n",
+    "for (a, b) in ((2, 5), (3, 3), (2, 8)):\n",
+    "    check(f\"expected live fraction closed form vs quadrature, Beta({a},{b}), G=16\", abs(live_fraction(a, b, 16) - live_quad(a, b, 16)), 1e-10)\n",
+    "a, b, G = 2.0, 8.0, 16\n",
+    "rho = live_fraction(a, b, G)\n",
+    "print(f\"   Beta({a:.0f},{b:.0f}) prompts (mean p={a/(a+b):.2f}), G={G}: live fraction rho = {rho:.3f}  ->  oversampling factor 1/rho = {1/rho:.2f}\")\n",
+    "# MC: prompts drawn until m live ones are found; expectation m / rho\n",
+    "torch.manual_seed(2)\n",
+    "gen = torch.Generator().manual_seed(2)\n",
+    "m_need, reps = 64, 600\n",
+    "draws = []\n",
+    "beta = torch.distributions.Beta(torch.tensor(a), torch.tensor(b))\n",
+    "for _ in range(reps):\n",
+    "    live = 0; drawn = 0\n",
+    "    while live < m_need:\n",
+    "        pp = beta.sample()\n",
+    "        k = torch.distributions.Binomial(G, pp).sample()\n",
+    "        drawn += 1; live += int(0 < k < G)\n",
+    "    draws.append(drawn)\n",
+    "d = torch.tensor(draws, dtype=torch.float64)\n",
+    "check(\"MC: prompts drawn to collect 64 live groups vs m/rho (z-score)\", abs(float(d.mean()) - m_need / rho) / (float(d.std()) / reps ** 0.5), 4.0, kind=\"mc\")\n",
+    "# Unbiasedness of dynamic sampling (exact, no MC): prompts q_j with success prob p_j and weights w_j; per-prompt Bernoulli-bandit gradient g = p_hat(1-p_hat)\n",
+    "# (centred advantage x score, see 0110).  Filtered batch: a prompt is accepted with probability P(live|q_j) (so accepted prompts have law w_j P(live_j)/rho) and its\n",
+    "# group contributes E[g | live, q_j] computed DIRECTLY from the binomial restricted to 0<k<G (not using g=0 on dead groups).\n",
+    "Gd = 16\n",
+    "pj = np.linspace(0.02, 0.98, 25); wj = pj ** (a - 1) * (1 - pj) ** (b - 1); wj /= wj.sum()\n",
+    "def binom_k(pp): return np.array([math.comb(Gd, k) * pp ** k * (1 - pp) ** (Gd - k) for k in range(Gd + 1)])\n",
+    "kk = np.arange(Gd + 1); gk = (kk / Gd) * (1 - kk / Gd)\n",
+    "Eg_all = np.array([float((binom_k(x) * gk).sum()) for x in pj])\n",
+    "Plive = np.array([float(binom_k(x)[1:Gd].sum()) for x in pj])\n",
+    "Eg_live = np.array([float((binom_k(x)[1:Gd] * gk[1:Gd]).sum()) / pl for x, pl in zip(pj, Plive)])\n",
+    "rho_j = float((wj * Plive).sum())\n",
+    "unfiltered = float((wj * Eg_all).sum())\n",
+    "filtered = float((wj * Plive / rho_j * Eg_live).sum())\n",
+    "check(\"dynamic sampling: E[g | accepted group] = (1/rho) * E[g | unfiltered group]  (exact, discrete prompt law)\", abs(filtered - unfiltered / rho_j), 1e-15)\n",
+    "# the dilution identity: dividing by ALL samples instead of live samples scales the gradient by exactly the live fraction\n",
+    "print(\"   with global token normalisation, dead groups stay in the denominator: signal scales by rho =\", round(rho, 3), \"(identity, no estimation)\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|mc|P(dead) MC vs p^G+(1-p)^G  (G=4, p=0.02)  |diff|/SE|9.828e-01|tol=4.0e+00|PASS\n",
+      "CHECK|mc|P(dead) MC vs p^G+(1-p)^G  (G=4, p=0.3)  |diff|/SE|6.955e-01|tol=4.0e+00|PASS\n",
+      "CHECK|mc|P(dead) MC vs p^G+(1-p)^G  (G=16, p=0.02)  |diff|/SE|4.989e-01|tol=4.0e+00|PASS\n",
+      "CHECK|mc|P(dead) MC vs p^G+(1-p)^G  (G=16, p=0.3)  |diff|/SE|2.052e+00|tol=4.0e+00|PASS\n",
+      "CHECK|exact|expected live fraction closed form vs quadrature, Beta(2,5), G=16|2.698e-14|tol=1.0e-10|PASS\n",
+      "CHECK|exact|expected live fraction closed form vs quadrature, Beta(3,3), G=16|2.798e-14|tol=1.0e-10|PASS\n",
+      "CHECK|exact|expected live fraction closed form vs quadrature, Beta(2,8), G=16|2.109e-14|tol=1.0e-10|PASS\n",
+      "   Beta(2,8) prompts (mean p=0.20), G=16: live fraction rho = 0.880  ->  oversampling factor 1/rho = 1.14\n",
+      "CHECK|mc|MC: prompts drawn to collect 64 live groups vs m/rho (z-score)|5.547e-01|tol=4.0e+00|PASS\n",
+      "CHECK|exact|dynamic sampling: E[g | accepted group] = (1/rho) * E[g | unfiltered group]  (exact, discrete prompt law)|0.000e+00|tol=1.0e-15|PASS\n",
+      "   with global token normalisation, dead groups stay in the denominator: signal scales by rho = 0.88 (identity, no estimation)\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Token-level loss: global denominator is a RATIO estimator (bias O(1/#samples)); constant denominator is exactly unbiased\n",
+    "Exact over all G-tuples (G=2,3) on the TinyLM."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "source": [
+    "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n",
+    "theta = lm.theta0.clone(); R = default_reward(lm); p = lm.probs(theta); S = score_matrix(lm, theta); N = lm.N\n",
+    "lens = lm.lengths\n",
+    "grad_J = (p[:, None] * (R - float(p @ R))[:, None] * S).sum(0)\n",
+    "Lbar = float((p * lens).sum())\n",
+    "def exact_tokenlevel(G, normaliser):\n",
+    "    \"\"\"E[ sum_i A_i * score(o_i) / D ] with centred advantages A_i = R_i - mean R; D = sum_i |o_i| (global) or G*Lbar (constant).\"\"\"\n",
+    "    if G == 2:\n",
+    "        R1, R2 = R[:, None].expand(N, N), R[None, :].expand(N, N)\n",
+    "        l1, l2 = lens[:, None].expand(N, N), lens[None, :].expand(N, N)\n",
+    "        P = p[:, None] * p[None, :]\n",
+    "        mu = (R1 + R2) / 2\n",
+    "        D = (l1 + l2) if normaliser == 'global' else torch.full((N, N), G * Lbar)\n",
+    "        w1 = P * (R1 - mu) / D; w2 = P * (R2 - mu) / D\n",
+    "        return w1.sum(1) @ S + w2.sum(0) @ S\n",
+    "    g = torch.zeros(S.shape[1])\n",
+    "    R2, R3 = R[:, None].expand(N, N), R[None, :].expand(N, N)\n",
+    "    l2, l3 = lens[:, None].expand(N, N), lens[None, :].expand(N, N)\n",
+    "    P23 = p[:, None] * p[None, :]\n",
+    "    for i in range(N):\n",
+    "        mu = (R[i] + R2 + R3) / 3\n",
+    "        D = (lens[i] + l2 + l3) if normaliser == 'global' else torch.full((N, N), G * Lbar)\n",
+    "        g += p[i] * ((P23 * (R[i] - mu) / D).sum()) * S[i]\n",
+    "        g += p[i] * ((P23 * (R2 - mu) / D).sum(1) @ S)\n",
+    "        g += p[i] * ((P23 * (R3 - mu) / D).sum(0) @ S)\n",
+    "    return g\n",
+    "for G in (2, 3):\n",
+    "    gc = exact_tokenlevel(G, 'const'); gg = exact_tokenlevel(G, 'global')\n",
+    "    check(f\"G={G}: constant denominator G*E|o| gives (1-1/G) grad J / (G E|o|) * G exactly\", maxabs(gc * Lbar, (1 - 1 / G) * grad_J), 1e-13)\n",
+    "    rel = float((gg - gc).norm() / gc.norm()); cs = float(gg @ grad_J / (gg.norm() * grad_J.norm()))\n",
+    "    print(f\"   G={G}: global-token-mean estimator: relative deviation from the unbiased one = {rel:.4f}, cosine to grad J = {cs:.5f}\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|G=2: constant denominator G*E|o| gives (1-1/G) grad J / (G E|o|) * G exactly|1.041e-17|tol=1.0e-13|PASS\n",
+      "   G=2: global-token-mean estimator: relative deviation from the unbiased one = 0.1591, cosine to grad J = 0.99588\n",
+      "CHECK|exact|G=3: constant denominator G*E|o| gives (1-1/G) grad J / (G E|o|) * G exactly|4.163e-17|tol=1.0e-13|PASS\n",
+      "   G=3: global-token-mean estimator: relative deviation from the unbiased one = 0.1000, cosine to grad J = 0.99779\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Overlong handling: soft punishment is a REWARD; loss masking is a reward equal to the baseline"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "source": [
+    "def soft_overlong(Ly, Lmax, Lcache):\n",
+    "    if Ly <= Lmax - Lcache: return 0.0\n",
+    "    if Ly <= Lmax: return ((Lmax - Lcache) - Ly) / Lcache\n",
+    "    return -1.0\n",
+    "Lmax, Lcache = 16384, 4096\n",
+    "check(\"soft overlong: 0 at Lmax-Lcache\", abs(soft_overlong(Lmax - Lcache, Lmax, Lcache)), 0.0)\n",
+    "check(\"soft overlong: -1 at Lmax (continuous with the truncation penalty)\", abs(soft_overlong(Lmax, Lmax, Lcache) + 1.0), 0.0)\n",
+    "ls = np.arange(1, 24000)\n",
+    "vals = np.array([soft_overlong(int(x), Lmax, Lcache) for x in ls])\n",
+    "slope = float(np.abs(np.diff(vals[:Lmax])).max())          # largest per-token change of the ramp (up to Lmax)\n",
+    "check(\"soft overlong: ramp slope is exactly 1/Lcache (|max slope * Lcache - 1|)\", abs(slope * Lcache - 1.0), 1e-9)\n",
+    "print(\"   hard truncation reward has a jump of size 1 at Lmax; the soft ramp spreads it over\", Lcache, \"tokens\")\n",
+    "# masking == baseline reward  (TinyLM: 'truncated' = hit the length limit L=4 without EOS)\n",
+    "trunc = (lm.mask.sum(1) == lm.L) & (lm.tokens[:, -1] != lm.EOS)\n",
+    "b = float(p @ R)\n",
+    "g_masked = (p[:, None] * ((~trunc).double() * (R - b))[:, None] * S).sum(0)               # loss masked on truncated samples\n",
+    "Rprime = torch.where(trunc, torch.full_like(R, b), R)                                      # truncated sample gets reward = baseline\n",
+    "g_rb = (p[:, None] * (Rprime - b)[:, None] * S).sum(0)\n",
+    "check(\"loss-masking == assigning the baseline value as reward to truncated samples (exact)\", maxabs(g_masked, g_rb), 1e-14)\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "Rlen = torch.where(trunc, torch.full_like(R, -1.0), torch.zeros_like(R)) + 0.0\n",
+    "g_total_reward = (p[:, None] * (R + Rlen - b)[:, None] * S).sum(0)\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "g_sep = torch.autograd.grad((lm.probs(th) * (R + Rlen)).sum(), th)[0]\n",
+    "check(\"shaping as reward: policy gradient of E[R + R_len] (separable objective)\", maxabs(g_total_reward, g_sep), 1e-14)\n",
+    "g_hard = (p[:, None] * (R - b)[:, None] * S).sum(0)                                        # truncated treated as plain failures (R=0)\n",
+    "print(f\"   ||g_masked - g_hard|| = {float((g_masked-g_hard).norm()):.4f}  vs ||grad J|| = {float(grad_J.norm()):.4f}  (filtering changes the objective)\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|soft overlong: 0 at Lmax-Lcache|0.000e+00|tol=0.0e+00|PASS\n",
+      "CHECK|exact|soft overlong: -1 at Lmax (continuous with the truncation penalty)|0.000e+00|tol=0.0e+00|PASS\n",
+      "CHECK|exact|soft overlong: ramp slope is exactly 1/Lcache (|max slope * Lcache - 1|)|0.000e+00|tol=1.0e-09|PASS\n",
+      "   hard truncation reward has a jump of size 1 at Lmax; the soft ramp spreads it over 4096 tokens\n",
+      "CHECK|exact|loss-masking == assigning the baseline value as reward to truncated samples (exact)|0.000e+00|tol=1.0e-14|PASS\n",
+      "CHECK|exact|shaping as reward: policy gradient of E[R + R_len] (separable objective)|2.776e-17|tol=1.0e-14|PASS\n",
+      "   ||g_masked - g_hard|| = 0.0124  vs ||grad J|| = 0.1825  (filtering changes the objective)\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Naive vs vectorised DAPO objective: asymmetric clip + global token-level normaliser + dynamic-sampling filter (gradient agreement)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "source": [
+    "th_old = lm.theta0.clone(); lq = lm.token_logprobs(th_old); q_old = lm.probs(th_old)\n",
+    "th_new = perturb(th_old, 1.0, seed=3)\n",
+    "gen = torch.Generator().manual_seed(17)\n",
+    "G, groups = 8, []\n",
+    "while len(groups) < 3:                                               # dynamic sampling: resample until the group is live (0 < k < G)\n",
+    "    idx = torch.multinomial(q_old, G, replacement=True, generator=gen)\n",
+    "    k = int(R[idx].sum())\n",
+    "    if 0 < k < G: groups.append(idx)\n",
+    "def dapo_naive(th):\n",
+    "    lp_all = lm.token_logprobs(th)\n",
+    "    tot, ntok = 0.0, 0\n",
+    "    for idx in groups:\n",
+    "        Rg = R[idx]; A = (Rg - Rg.mean()) / (Rg.std(unbiased=False) + 1e-6)\n",
+    "        for j, i in enumerate(idx.tolist()):\n",
+    "            for t in range(int(lm.lengths[i])):\n",
+    "                r = torch.exp(lp_all[i, t] - lq[i, t])\n",
+    "                a = A[j]\n",
+    "                term = a * torch.clamp(r, max=1 + eh) if a > 0 else a * torch.clamp(r, min=1 - el)\n",
+    "                tot = tot + term; ntok += 1\n",
+    "    return -tot / ntok\n",
+    "def dapo_vec(th):\n",
+    "    lp_all = lm.token_logprobs(th)\n",
+    "    idx = torch.cat(groups); Rg = R[idx].reshape(len(groups), G)\n",
+    "    A = ((Rg - Rg.mean(1, keepdim=True)) / (Rg.std(1, unbiased=False, keepdim=True) + 1e-6)).reshape(-1)\n",
+    "    r = (lp_all[idx] - lq[idx]).exp(); mk = lm.mask[idx]\n",
+    "    f = torch.where(A[:, None] > 0, A[:, None] * torch.clamp(r, max=1 + eh), A[:, None] * torch.clamp(r, min=1 - el))\n",
+    "    return -(f * mk).sum() / mk.sum()\n",
+    "t1 = th_new.clone().requires_grad_(True); t2 = th_new.clone().requires_grad_(True)\n",
+    "check(\"naive triple loop == vectorised DAPO loss (value)\", abs(float(dapo_naive(t1).detach()) - float(dapo_vec(t2).detach())), 1e-13)\n",
+    "g1 = torch.autograd.grad(dapo_naive(t1), t1)[0]; g2 = torch.autograd.grad(dapo_vec(t2), t2)[0]\n",
+    "check(\"naive triple loop == vectorised DAPO loss (gradient)\", maxabs(g1, g2), 1e-13)\n",
+    "print(\"done; checks:\", len(CHECKS))"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|naive triple loop == vectorised DAPO loss (value)|1.665e-16|tol=1.0e-13|PASS\n",
+      "CHECK|exact|naive triple loop == vectorised DAPO loss (gradient)|0.000e+00|tol=1.0e-13|PASS\n",
+      "done; checks: 24\n"
+     ]
+    }
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

code/notebooks/0107-cispo.ipynb

@@ -0,0 +1,596 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 0107 - CISPO (MiniMax-M1 Sec. 3.1): clip the importance weight, not the update\n",
+    "Companion to `notes/0107 - CISPO.pdf`.  fp64; exact enumeration on the 121-sequence toy LM."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "source": [
+    "import sys, math\n",
+    "sys.path.insert(0, \"..\")\n",
+    "from _m1_lib import *\n",
+    "import numpy as np\n",
+    "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n",
+    "R = default_reward(lm)\n",
+    "th_old = lm.theta0.clone(); q = lm.probs(th_old); N, L = lm.N, lm.L\n",
+    "b = float(q @ R)\n",
+    "A_tok = (R - b)[:, None] * lm.mask                       # outcome advantage, same for all tokens of a sequence\n",
+    "lq = lm.token_logprobs(th_old)\n",
+    "Ttens = token_score_tensor(lm, th_old)                    # score of each token at THETA_OLD (used to build gradients analytically)\n",
+    "eh, el = 0.28, 0.2"
+   ],
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Gradient weights.  Every method has gradient  sum_t  w(r_t, A_t) * A_t * grad log pi_theta(o_t | s_t)\n",
+    "  PPO/GRPO (clip lo/hi):   w = r * 1[live]       live = (A>0 & r<1+eh) | (A<0 & r>1-el)\n",
+    "  CISPO:                   w = min(r, 1+eh)      (A-independent, always live)      [loss = sg(min(r,1+eh)) * A * log pi]\n",
+    "  token IS-PG (no clip):   w = r\n",
+    "We verify against autograd of the actual loss functions on the toy model at a perturbed theta."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "source": [
+    "def ratio_and_lp(th):\n",
+    "    lp = lm.token_logprobs(th)\n",
+    "    r = (lp - lq).exp() * lm.mask + (1 - lm.mask)\n",
+    "    return r, lp\n",
+    "def loss_ppo(th):\n",
+    "    r, lp = ratio_and_lp(th)\n",
+    "    f = torch.where(A_tok > 0, A_tok * torch.clamp(r, max=1 + eh), A_tok * torch.clamp(r, min=1 - el))\n",
+    "    return -(q[:, None] * f * lm.mask).sum()\n",
+    "def loss_cispo(th, eps_hi=eh):\n",
+    "    r, lp = ratio_and_lp(th)\n",
+    "    rhat = torch.clamp(r, max=1 + eps_hi).detach()                    # sg( clip(r) )\n",
+    "    return -(q[:, None] * rhat * A_tok * lp * lm.mask).sum()\n",
+    "def loss_isPG(th):\n",
+    "    r, lp = ratio_and_lp(th)\n",
+    "    return -(q[:, None] * r.detach() * A_tok * lp * lm.mask).sum()\n",
+    "def weights(th, method):\n",
+    "    r, _ = ratio_and_lp(th); r = r.detach()\n",
+    "    if method == 'ppo': live = ((A_tok > 0) & (r < 1 + eh)) | ((A_tok < 0) & (r > 1 - el)); return r * live\n",
+    "    if method == 'cispo': return torch.clamp(r, max=1 + eh)\n",
+    "    if method == 'is': return r\n",
+    "def grad_from_weights(th, w):\n",
+    "    # sum_{o,t} q(o) w A  grad log pi_theta(o_t|s_t)  with the score evaluated at theta (autograd through lp)\n",
+    "    t = th.clone().requires_grad_(True)\n",
+    "    lp = lm.token_logprobs(t)\n",
+    "    return torch.autograd.grad((q[:, None] * w * A_tok * lp * lm.mask).sum(), t)[0]\n",
+    "th = perturb(th_old, 1.5, seed=11)\n",
+    "r_now, _ = ratio_and_lp(th)\n",
+    "print(\"ratio range over (o,t):  min %.3f  max %.3f\" % (float(r_now[lm.mask > 0].min()), float(r_now[lm.mask > 0].max())))\n",
+    "for name, lossf in ((\"ppo\", loss_ppo), (\"cispo\", loss_cispo), (\"is\", loss_isPG)):\n",
+    "    t = th.clone().requires_grad_(True)\n",
+    "    g_auto = torch.autograd.grad(lossf(t), t)[0]\n",
+    "    g_w = -grad_from_weights(th, weights(th, name))\n",
+    "    check(f\"{name}: autograd of the loss == sum_t w(r,A) A grad log pi  (weights table)\", maxabs(g_auto, g_w), 1e-13)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "ratio range over (o,t):  min 0.455  max 1.992\n",
+      "CHECK|exact|ppo: autograd of the loss == sum_t w(r,A) A grad log pi  (weights table)|1.388e-17|tol=1.0e-13|PASS\n",
+      "CHECK|exact|cispo: autograd of the loss == sum_t w(r,A) A grad log pi  (weights table)|0.000e+00|tol=1.0e-13|PASS\n",
+      "CHECK|exact|is: autograd of the loss == sum_t w(r,A) A grad log pi  (weights table)|0.000e+00|tol=1.0e-13|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Exactly which tokens does PPO/GRPO silence and CISPO keep?"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "source": [
+    "live = ((A_tok > 0) & (r_now < 1 + eh)) | ((A_tok < 0) & (r_now > 1 - el))\n",
+    "sel = lm.mask > 0\n",
+    "pw = q[:, None] * lm.mask                                  # probability weight of each (o,t) token occurrence\n",
+    "tot = float(pw.sum())\n",
+    "zeroed = sel & ~live\n",
+    "cells = {\"A>0, r>1+eh  (silenced by PPO)\": sel & (A_tok > 0) & (r_now > 1 + eh),\n",
+    "         \"A<0, r<1-el  (silenced by PPO)\": sel & (A_tok < 0) & (r_now < 1 - el),\n",
+    "         \"A<0, r>1+eh  (PPO weight r unbounded; CISPO caps at 1+eh)\": sel & (A_tok < 0) & (r_now > 1 + eh)}\n",
+    "for k, m in cells.items(): print(f\"   {k:60s}: {int(m.sum()):4d} token slots, probability mass {float((pw*m).sum())/tot:.4f}\")\n",
+    "Wc = weights(th, 'cispo'); Wp = weights(th, 'ppo')\n",
+    "silenced_by_ppo_kept_by_cispo = sel & (Wp == 0) & (Wc > 0)\n",
+    "check(\"{tokens with PPO weight 0 and CISPO weight>0} == cells 1 U 2 exactly (symmetric difference size)\", float((silenced_by_ppo_kept_by_cispo ^ (cells[\"A>0, r>1+eh  (silenced by PPO)\"] | cells[\"A<0, r<1-el  (silenced by PPO)\"])).sum()), 0.0)\n",
+    "check(\"CISPO never silences a token (min weight over all (o,t) with A != 0)\", abs(float(Wc[sel].min()) - float(r_now[sel].min())) if float(r_now[sel].min()) < 1 + eh else 0.0, 1e-15)\n",
+    "check(\"PPO silenced set == the MiniMax-M1 mask M_{i,t}=0 of Eq. 7 (symmetric difference size)\",\n",
+    "      float((zeroed ^ (((A_tok > 0) & (r_now > 1 + eh)) | ((A_tok < 0) & (r_now < 1 - el))) & sel).sum()), 0.0)\n",
+    "# weight table\n",
+    "print(\"\\n   gradient weight w(r,A)  (loss gradient = w * A * grad log pi),  eps_hi=0.28, eps_lo=0.2\")\n",
+    "print(\"      r      PPO(A>0)  PPO(A<0)   CISPO   unclipped-IS\")\n",
+    "for r in (0.3, 0.8, 1.0, 1.28, 2.0, 5.0, 20.0):\n",
+    "    wp_pos = r if r < 1 + eh else 0.0; wp_neg = r if r > 1 - el else 0.0\n",
+    "    print(f\"    {r:5.2f}   {wp_pos:7.2f}   {wp_neg:7.2f}   {min(r, 1+eh):6.2f}   {r:7.2f}\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "   A>0, r>1+eh  (silenced by PPO)                              :    7 token slots, probability mass 0.0134\n",
+      "   A<0, r<1-el  (silenced by PPO)                              :   63 token slots, probability mass 0.0551\n",
+      "   A<0, r>1+eh  (PPO weight r unbounded; CISPO caps at 1+eh)   :   83 token slots, probability mass 0.0160\n",
+      "CHECK|exact|{tokens with PPO weight 0 and CISPO weight>0} == cells 1 U 2 exactly (symmetric difference size)|0.000e+00|tol=0.0e+00|PASS\n",
+      "CHECK|exact|CISPO never silences a token (min weight over all (o,t) with A != 0)|0.000e+00|tol=1.0e-15|PASS\n",
+      "CHECK|exact|PPO silenced set == the MiniMax-M1 mask M_{i,t}=0 of Eq. 7 (symmetric difference size)|0.000e+00|tol=0.0e+00|PASS\n",
+      "\n",
+      "   gradient weight w(r,A)  (loss gradient = w * A * grad log pi),  eps_hi=0.28, eps_lo=0.2\n",
+      "      r      PPO(A>0)  PPO(A<0)   CISPO   unclipped-IS\n",
+      "     0.30      0.30      0.00     0.30      0.30\n",
+      "     0.80      0.80      0.00     0.80      0.80\n",
+      "     1.00      1.00      1.00     1.00      1.00\n",
+      "     1.28      0.00      1.28     1.28      1.28\n",
+      "     2.00      0.00      2.00     1.28      2.00\n",
+      "     5.00      0.00      5.00     1.28      5.00\n",
+      "    20.00      0.00     20.00     1.28     20.00\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Unified mask formulation (paper Eq. 6-7):  J = sg(r_hat) A log pi * M ;  with r_hat = r it reproduces PPO's gradient"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "source": [
+    "def loss_unify(th, rhat_clip=None):\n",
+    "    r, lp = ratio_and_lp(th)\n",
+    "    M = (~(((A_tok > 0) & (r > 1 + eh)) | ((A_tok < 0) & (r < 1 - el)))).double()\n",
+    "    rhat = r if rhat_clip is None else torch.clamp(r, max=rhat_clip)\n",
+    "    return -(q[:, None] * rhat.detach() * A_tok * lp * M * lm.mask).sum()\n",
+    "t = th.clone().requires_grad_(True); g_u = torch.autograd.grad(loss_unify(t), t)[0]\n",
+    "t = th.clone().requires_grad_(True); g_p = torch.autograd.grad(loss_ppo(t), t)[0]\n",
+    "check(\"unified formulation with r_hat=r and PPO mask == PPO gradient\", maxabs(g_u, g_p), 1e-13)\n",
+    "# without clipping CISPO reduces to the standard (token-level IS) policy gradient; at theta_old it is the on-policy gradient\n",
+    "t = th.clone().requires_grad_(True); g_c_inf = torch.autograd.grad(loss_cispo(t, eps_hi=1e9), t)[0]\n",
+    "t = th.clone().requires_grad_(True); g_is = torch.autograd.grad(loss_isPG(t), t)[0]\n",
+    "check(\"CISPO with eps_hi=inf == token-level IS policy gradient\", maxabs(g_c_inf, g_is), 1e-14)\n",
+    "t = th_old.clone().requires_grad_(True); g_c0 = torch.autograd.grad(loss_cispo(t), t)[0]\n",
+    "t2 = th_old.clone().requires_grad_(True); gJ = torch.autograd.grad((lm.probs(t2) * R).sum(), t2)[0]\n",
+    "check(\"at theta_old CISPO == exact policy gradient grad J (r=1, nothing clipped)\", maxabs(-g_c0, gJ), 1e-14)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|unified formulation with r_hat=r and PPO mask == PPO gradient|1.388e-17|tol=1.0e-13|PASS\n",
+      "CHECK|exact|CISPO with eps_hi=inf == token-level IS policy gradient|0.000e+00|tol=1.0e-14|PASS\n",
+      "CHECK|exact|at theta_old CISPO == exact policy gradient grad J (r=1, nothing clipped)|4.163e-17|tol=1.0e-14|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Bias of CISPO relative to the (token-level) IS gradient, and the bound  ||bias|| <= sum_t E[ (r - 1 - eps)_+ |A| ||score|| ]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "source": [
+    "def cispo_bias(th, eps):\n",
+    "    r, _ = ratio_and_lp(th); r = r.detach()\n",
+    "    Tth = token_score_tensor(lm, th)                                   # score at theta\n",
+    "    diff = (torch.clamp(r, max=1 + eps) - r)                            # r_hat - r  (<= 0)\n",
+    "    return (q[:, None, None] * (diff * A_tok)[:, :, None] * Tth).sum((0, 1)), Tth\n",
+    "for eps_ in (0.05, 0.1, 0.28):\n",
+    "    bias, Tth = cispo_bias(th, eps_)\n",
+    "    r, _ = ratio_and_lp(th)\n",
+    "    bound = float((q[:, None] * (r - 1 - eps_).clamp(min=0) * A_tok.abs() * Tth.norm(dim=2) * lm.mask).sum())\n",
+    "    print(f\"   eps_hi={eps_:4.2f}:  ||bias|| = {float(bias.norm()):.4e}   bound = {bound:.4e}   ||token-IS gradient|| = {float((q[:,None,None]*(r*A_tok)[:,:,None]*Tth).sum((0,1)).norm()):.4e}\")\n",
+    "    assert float(bias.norm()) <= bound + 1e-15\n",
+    "check(\"bias bound holds (3 values of eps_hi)\", 0.0, 0.0)\n",
+    "# for comparison: how far each method's gradient is from the token-level IS gradient\n",
+    "t_ = th.clone().requires_grad_(True); gI = torch.autograd.grad(loss_isPG(t_), t_)[0]\n",
+    "t_ = th.clone().requires_grad_(True); gP = torch.autograd.grad(loss_ppo(t_), t_)[0]\n",
+    "t_ = th.clone().requires_grad_(True); gC = torch.autograd.grad(loss_cispo(t_), t_)[0]\n",
+    "print(f\"   ||g_PPO - g_IS|| = {float((gP-gI).norm()):.4e}   ||g_CISPO - g_IS|| = {float((gC-gI).norm()):.4e}   (eps_hi=0.28, eps_lo=0.2; ||g_IS|| = {float(gI.norm()):.4e})\")\n",
+    "# distribution-free per-state bound  E_old[(r-1-eps)_+] = sum_a (pi - (1+eps) pi_old)_+  <=  TV(pi, pi_old)\n",
+    "worst = 0.0\n",
+    "for s in lm.prefixes():\n",
+    "    po, pn = lm.cond_probs(th_old, s), lm.cond_probs(th, s)\n",
+    "    lhs = float(((pn - 1.28 * po).clamp(min=0)).sum()); tv = 0.5 * float((pn - po).abs().sum())\n",
+    "    worst = max(worst, lhs - tv)\n",
+    "check(\"per-state E_old[(r-1-eps)_+] <= TV(pi,pi_old)  (max violation, must be <= 0)\", max(0.0, worst), 0.0)\n",
+    "# variance: gradient weights are bounded for CISPO, unbounded for PPO's live A<0 tokens\n",
+    "r, _ = ratio_and_lp(th)\n",
+    "print(f\"   max gradient weight over (o,t):  PPO {float(weights(th,'ppo').max()):.3f}   CISPO {float(weights(th,'cispo').max()):.3f}   unclipped IS {float(weights(th,'is').max()):.3f}\")\n",
+    "assert float(weights(th, 'cispo').max()) <= 1 + eh + 1e-15\n",
+    "check(\"CISPO weights bounded by 1+eps_hi\", max(0.0, float(weights(th, 'cispo').max()) - (1 + eh)), 0.0)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "   eps_hi=0.05:  ||bias|| = 1.3732e-02   bound = 3.7511e-02   ||token-IS gradient|| = 1.7479e-01\n",
+      "   eps_hi=0.10:  ||bias|| = 9.2089e-03   bound = 2.3875e-02   ||token-IS gradient|| = 1.7479e-01\n",
+      "   eps_hi=0.28:  ||bias|| = 1.0099e-03   bound = 4.1370e-03   ||token-IS gradient|| = 1.7479e-01\n",
+      "CHECK|exact|bias bound holds (3 values of eps_hi)|0.000e+00|tol=0.0e+00|PASS\n",
+      "   ||g_PPO - g_IS|| = 2.2235e-02   ||g_CISPO - g_IS|| = 1.0099e-03   (eps_hi=0.28, eps_lo=0.2; ||g_IS|| = 1.7479e-01)\n",
+      "CHECK|exact|per-state E_old[(r-1-eps)_+] <= TV(pi,pi_old)  (max violation, must be <= 0)|0.000e+00|tol=0.0e+00|PASS\n",
+      "   max gradient weight over (o,t):  PPO 1.992   CISPO 1.280   unclipped IS 1.992\n",
+      "CHECK|exact|CISPO weights bounded by 1+eps_hi|0.000e+00|tol=0.0e+00|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Stop-gradient semantics: three look-alike implementations, three different gradients\n",
+    "  (a) correct:    loss = - sg(min(r,1+eh)) * A * log pi         -> grad = -A*min(r,1+eh) * grad log pi\n",
+    "  (b) forgot sg:  loss = - min(r,1+eh) * A * log pi             -> on the UNclipped set grad = -A * r * (1 + log pi) * grad log pi : an extra  A r log pi  term\n",
+    "  (c) ratio form: loss = - min(r,1+eh) * A  (PPO-like)           -> grad = -A * r * 1[r<1+eh] * grad log pi   (zero beyond the clip!)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "source": [
+    "lp_old = torch.tensor(-2.3)\n",
+    "def three(lp_new, A, eps=0.28):\n",
+    "    lp = torch.tensor(lp_new, requires_grad=True); r = (lp - lp_old).exp()\n",
+    "    ga = torch.autograd.grad(-(torch.clamp(r, max=1 + eps).detach() * A * lp), lp, retain_graph=True)[0]\n",
+    "    gb = torch.autograd.grad(-(torch.clamp(r, max=1 + eps) * A * lp), lp, retain_graph=True)[0]\n",
+    "    gc = torch.autograd.grad(-(torch.clamp(r, max=1 + eps) * A), lp)[0]\n",
+    "    return float(ga), float(gb), float(gc), float(r.detach())\n",
+    "print(\"   lp_old=-2.3;   gradient wrt log pi_theta\")\n",
+    "print(\"   lp_new     r     A     (a) CISPO     (b) no sg     (c) ratio-clip\")\n",
+    "for lpn, A in ((-2.3, 1.0), (-1.5, 1.0), (-1.5, -1.0), (-3.0, 1.0), (-0.5, 1.0), (-0.5, -1.0)):\n",
+    "    ga, gb, gc, rr = three(lpn, A)\n",
+    "    print(f\"   {lpn:6.2f}  {rr:6.2f}  {A:+.0f}   {ga:+9.3f}   {gb:+9.3f}   {gc:+9.3f}\")\n",
+    "    expect_a = -A * min(rr, 1.28)\n",
+    "    assert abs(ga - expect_a) < 1e-12\n",
+    "    expect_c = -A * rr * (rr < 1.28)\n",
+    "    assert abs(gc - expect_c) < 1e-12\n",
+    "    expect_b = -A * rr * (1 + lpn) if rr < 1.28 else -A * 1.28           # clamp kills the extra term beyond the clip\n",
+    "    assert abs(gb - expect_b) < 1e-12\n",
+    "check(\"(a) = -A min(r,1+eh);  (b) = -A r (1+log pi) on the live set;  (c) = -A r 1[r<1+eh]   (6 test points, exact)\", 0.0, 0.0)\n",
+    "ga, gb, gc, rr = three(-2.3, 1.0)\n",
+    "print(f\"   at r=1 (inside the clip), A=+1, log pi=-2.3: (a)={ga:+.3f}   (b)={gb:+.3f}  <- sign flipped by the spurious log-pi factor (1+log pi = -1.3)\")\n",
+    "assert ga < 0 < gb\n",
+    "check(\"without sg a live token's gradient weight becomes r(1+log pi): wrong sign when pi < 1/e\", 0.0, 0.0)\n",
+    "ga, gb, gc, rr = three(-1.5, 1.0)\n",
+    "print(f\"   at r={rr:.2f} > 1+eps: (a)={ga:+.3f} capped weight kept;  (b)={gb:+.3f} identical (clamp zeroes the extra term);  (c)={gc:+.3f} silenced\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "   lp_old=-2.3;   gradient wrt log pi_theta\n",
+      "   lp_new     r     A     (a) CISPO     (b) no sg     (c) ratio-clip\n",
+      "    -2.30    1.00  +1      -1.000      +1.300      -1.000\n",
+      "    -1.50    2.23  +1      -1.280      -1.280      +0.000\n",
+      "    -1.50    2.23  -1      +1.280      +1.280      +0.000\n",
+      "    -3.00    0.50  +1      -0.497      +0.993      -0.497\n",
+      "    -0.50    6.05  +1      -1.280      -1.280      +0.000\n",
+      "    -0.50    6.05  -1      +1.280      +1.280      +0.000\n",
+      "CHECK|exact|(a) = -A min(r,1+eh);  (b) = -A r (1+log pi) on the live set;  (c) = -A r 1[r<1+eh]   (6 test points, exact)|0.000e+00|tol=0.0e+00|PASS\n",
+      "   at r=1 (inside the clip), A=+1, log pi=-2.3: (a)=-1.000   (b)=+1.300  <- sign flipped by the spurious log-pi factor (1+log pi = -1.3)\n",
+      "CHECK|exact|without sg a live token's gradient weight becomes r(1+log pi): wrong sign when pi < 1/e|0.000e+00|tol=0.0e+00|PASS\n",
+      "   at r=2.23 > 1+eps: (a)=-1.280 capped weight kept;  (b)=-1.280 identical (clamp zeroes the extra term);  (c)=+0.000 silenced\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. The 'fork token' mechanism over K = 16 off-policy updates on one batch (MiniMax-M1 setting)\n",
+    "One decision with three tokens: two common tokens and a RARE token w (pi_old = 0.01) with positive advantage.\n",
+    "Exact expectation over the old distribution; plain SGD on the logits; K = 16 inner steps."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "source": [
+    "def fork(method, K=16, lr=30.0, eh_=0.28, el_=0.2, p_w=0.01):\n",
+    "    pi_old = torch.tensor([0.60, 0.39 - 0.0, p_w]); pi_old = pi_old / pi_old.sum()\n",
+    "    reward = torch.tensor([0.0, 0.0, 1.0]); A = reward - float(pi_old @ reward)       # A_w > 0, A_common < 0\n",
+    "    z = pi_old.log().clone().requires_grad_(True)\n",
+    "    traj, gw = [], []\n",
+    "    for k in range(K):\n",
+    "        pi = torch.softmax(z, 0); r = (pi / pi_old)\n",
+    "        if method == 'ppo':\n",
+    "            f = torch.where(A > 0, A * torch.clamp(r, max=1 + eh_), A * torch.clamp(r, min=1 - el_))\n",
+    "            obj = (pi_old * f).sum()\n",
+    "        else:\n",
+    "            rhat = torch.clamp(r, max=1 + eh_).detach()\n",
+    "            obj = (pi_old * rhat * A * pi.log()).sum()\n",
+    "        g = torch.autograd.grad(obj, z)[0]\n",
+    "        # direct contribution of token w's OWN term to the gradient on its own logit (diagnostic)\n",
+    "        rd = r.detach(); wt = float((1 + eh_) if (method == 'cispo' and rd[2] > 1 + eh_) else (rd[2] if (method == 'cispo' or rd[2] < 1 + eh_) else 0.0))\n",
+    "        gw.append(wt)\n",
+    "        with torch.no_grad(): z += lr * g\n",
+    "        traj.append(float(torch.softmax(z.detach(), 0)[2]))\n",
+    "    return traj, gw, float(r[2].detach())\n",
+    "tp, gwp, rp = fork('ppo'); tc, gwc, rc = fork('cispo')\n",
+    "print(\"   after update #\" + \" \".join(f\"{k+1:6d}\" for k in range(0, 16, 3)))\n",
+    "print(\"   p_w PPO:   \" + \" \".join(f\"{tp[k]:6.3f}\" for k in range(0, 16, 3)))\n",
+    "print(\"   p_w CISPO: \" + \" \".join(f\"{tc[k]:6.3f}\" for k in range(0, 16, 3)))\n",
+    "print(\"   weight on w's own term PPO:   \" + \" \".join(f\"{gwp[k]:6.2f}\" for k in range(0, 16, 3)))\n",
+    "print(\"   weight on w's own term CISPO: \" + \" \".join(f\"{gwc[k]:6.2f}\" for k in range(0, 16, 3)))\n",
+    "n_zero_ppo = sum(1 for x in gwp if x == 0.0)\n",
+    "print(f\"   PPO: token w gets zero direct gradient in {n_zero_ppo} of 16 inner steps;  CISPO: {sum(1 for x in gwc if x==0.0)} of 16\")\n",
+    "assert n_zero_ppo >= 8 and all(x > 0 for x in gwc)\n",
+    "check(\"PPO silences the rare token for >= 8 of 16 inner steps, CISPO for none\", 0.0, 0.0)\n",
+    "print(f\"   final p_w after 16 updates: PPO {tp[-1]:.3f}   CISPO {tc[-1]:.3f}\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "   after update #     1      4      7     10     13     16\n",
+      "   p_w PPO:    0.016  0.016  0.016  0.017  0.017  0.017\n",
+      "   p_w CISPO:  0.016  0.081  0.310  0.624  0.790  0.862\n",
+      "   weight on w's own term PPO:     1.00   0.00   0.00   0.00   0.00   0.00\n",
+      "   weight on w's own term CISPO:   1.00   1.28   1.28   1.28   1.28   1.28\n",
+      "   PPO: token w gets zero direct gradient in 15 of 16 inner steps;  CISPO: 0 of 16\n",
+      "CHECK|exact|PPO silences the rare token for >= 8 of 16 inner steps, CISPO for none|0.000e+00|tol=0.0e+00|PASS\n",
+      "   final p_w after 16 updates: PPO 0.017   CISPO 0.862\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7. Naive vs efficient implementation of the CISPO loss, token-level normalisation by total tokens"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "source": [
+    "Gs = 8\n",
+    "gen = torch.Generator().manual_seed(5)\n",
+    "idx = torch.multinomial(q, Gs, replacement=True, generator=gen)\n",
+    "Rg = R[idx]; Ag = Rg - Rg.mean()\n",
+    "t = perturb(th_old, 0.8, seed=2).clone().requires_grad_(True)\n",
+    "lp = lm.token_logprobs(t)[idx]; lq_ = lq[idx]; mk = lm.mask[idx]\n",
+    "def cispo_naive(lp, lq_):\n",
+    "    tot, n = 0.0, 0\n",
+    "    for i in range(Gs):\n",
+    "        for tt in range(int(mk[i].sum())):\n",
+    "            r = math.exp(float((lp[i, tt] - lq_[i, tt]).detach()))\n",
+    "            tot = tot + min(r, 1 + eh) * Ag[i] * lp[i, tt]; n += 1\n",
+    "    return -tot / n\n",
+    "def cispo_vec(lp, lq_):\n",
+    "    r = (lp - lq_).exp().detach()\n",
+    "    return -((torch.clamp(r, max=1 + eh) * Ag[:, None] * lp * mk).sum() / mk.sum())\n",
+    "gn = torch.autograd.grad(cispo_naive(lp, lq_), t, retain_graph=True)[0]\n",
+    "gv = torch.autograd.grad(cispo_vec(lp, lq_), t, retain_graph=True)[0]\n",
+    "check(\"naive per-token loop == vectorised masked CISPO loss (gradient)\", maxabs(gn, gv), 1e-14)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|naive per-token loop == vectorised masked CISPO loss (gradient)|0.000e+00|tol=1.0e-14|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 8. Leanstral variant (Mistral technical report, Sec. 2.3): sampler ratio, masking of tool/prompt tokens, outcome broadcast\n",
+    "Loss there:  `L = -(1/sum_i|o_i|) sum_i sum_t sg(min(pi_theta/pi_sample, c)) * A_i * log pi_theta(a_it | s_it, T_i)`, `T_i` = tool calls and responses used as context only,\n",
+    "one verifier outcome for every model token.  Toy: two-turn tabular agent (4 actions per turn), turn-2 context = turn-1 action, feedback token = a fixed function of turn 1 (not scored).\n",
+    "`R ~ Bern(T[y1, y2])`.  Everything below is enumerated exactly over the 32 episodes (y1, y2, R)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "source": [
+    "import numpy as _np\n",
+    "_rng = _np.random.default_rng(123)\n",
+    "_A = 4\n",
+    "_T = _rng.uniform(0.05, 0.95, (_A, _A))\n",
+    "def _sm(x): e = _np.exp(x - x.max(-1, keepdims=True)); return e / e.sum(-1, keepdims=True)\n",
+    "def _unpack(th): return th[:_A], th[_A:].reshape(_A, _A)\n",
+    "th_cur = _rng.normal(0, 0.8, _A + _A * _A)\n",
+    "th_smp = th_cur + _rng.normal(0, 0.35, _A + _A * _A)              # sampler = lagged + numerically perturbed copy\n",
+    "def _probs(th): a, b = _unpack(th); return _sm(a), _sm(b)\n",
+    "def _J(th): p1, p2 = _probs(th); return float((p1[:, None] * p2 * _T).sum())\n",
+    "def _score(th, y1, y2):                                           # (s1, s2): grads of log pi(y1), log pi(y2|y1) wrt the 20 logits\n",
+    "    p1, p2 = _probs(th); s1 = _np.zeros(_A + _A * _A); s2 = _np.zeros(_A + _A * _A)\n",
+    "    s1[:_A] = _np.eye(_A)[y1] - p1\n",
+    "    blk = _np.eye(_A)[y2] - p2[y1]; s2[_A + y1 * _A: _A + (y1 + 1) * _A] = blk\n",
+    "    return s1, s2\n",
+    "def _episodes(th_sample):                                         # list of (prob under sampler, y1, y2, R)\n",
+    "    p1, p2 = _probs(th_sample); out = []\n",
+    "    for y1 in range(_A):\n",
+    "        for y2 in range(_A):\n",
+    "            for r in (0, 1):\n",
+    "                out.append((p1[y1] * p2[y1, y2] * (_T[y1, y2] if r else 1 - _T[y1, y2]), y1, y2, r))\n",
+    "    return out\n",
+    "def _fd_grad(th, h=1e-6):\n",
+    "    g = _np.zeros_like(th)\n",
+    "    for i in range(len(th)):\n",
+    "        e = _np.zeros_like(th); e[i] = h; g[i] = (_J(th + e) - _J(th - e)) / (2 * h)\n",
+    "    return g\n",
+    "g_true = _fd_grad(th_cur)\n",
+    "pc1, pc2 = _probs(th_cur); ps1, ps2 = _probs(th_smp)\n",
+    "def _rho(y1, y2): return pc1[y1] / ps1[y1], pc2[y1, y2] / ps2[y1, y2]\n",
+    "# (a) on-policy score-function identity\n",
+    "g_on = sum(p * r * sum(_score(th_cur, y1, y2)) for p, y1, y2, r in _episodes(th_cur))\n",
+    "check(\"8a. E_pi[R (s1+s2)] = finite-difference grad J\", float(_np.abs(g_on - g_true).max()), 1e-8)\n",
+    "# (b) sampler episodes, full sequence ratio rho1*rho2: exact\n",
+    "g_seq = sum(p * _rho(y1, y2)[0] * _rho(y1, y2)[1] * r * sum(_score(th_cur, y1, y2)) for p, y1, y2, r in _episodes(th_smp))\n",
+    "check(\"8b. E_mu[rho1 rho2 R (s1+s2)] = grad J exactly (sampler ratio, full sequence weight)\", float(_np.abs(g_seq - g_true).max()), 1e-9)\n",
+    "# (c) token-level ratios only (what the objective uses): biased by the dropped prefix ratio; vanishes when sampler = policy\n",
+    "def _tok_grad(th_sample, clip=None, baseline=0.0):\n",
+    "    g = 0\n",
+    "    for p, y1, y2, r in _episodes(th_sample):\n",
+    "        q1, q2 = _rho(y1, y2) if th_sample is th_smp else (1.0, 1.0)\n",
+    "        if clip is not None: q1, q2 = min(q1, clip), min(q2, clip)\n",
+    "        s1, s2 = _score(th_cur, y1, y2); g = g + p * (r - baseline) * (q1 * s1 + q2 * s2)\n",
+    "    return g\n",
+    "g_tok = _tok_grad(th_smp)\n",
+    "bias_tok = float(_np.linalg.norm(g_tok - g_true)); gn = float(_np.linalg.norm(g_true))\n",
+    "print(f\"   token-level sampler ratio, unclipped: ||bias|| = {bias_tok:.3e} (||grad J|| = {gn:.3e})\")\n",
+    "check(\"8c. token-level weights are biased for a mismatched sampler (> 1e-4)\", 0.0 if bias_tok > 1e-4 else 1.0, 0.0)\n",
+    "th_save = th_smp; th_smp = th_cur.copy()                          # sampler == policy: weights are 1\n",
+    "check(\"8c. ... and exact when the sampler equals the policy\", float(_np.abs(sum(p * r * sum(_score(th_cur, y1, y2)) for p, y1, y2, r in _episodes(th_cur)) - g_true).max()), 1e-8)\n",
+    "th_smp = th_save\n",
+    "# (d) a constant baseline does not bias the token-level estimator, even under mismatch (E_mu[rho_t s_t | context] = 0)\n",
+    "b0 = _J(th_cur)\n",
+    "g_b = _tok_grad(th_smp, baseline=b0)\n",
+    "db = _np.abs(g_tok - g_b - 0).max()\n",
+    "check(\"8d. constant baseline b adds exactly 0 to the token-level sampler-ratio gradient: max |E_mu[sum_t rho_t b s_t]| over 20 coordinates\", float(db), 1e-14)\n",
+    "# (e) clip bias bound with the sampler as the base measure: sum_a (pi - c mu)_+ <= TV(pi, mu) at every state\n",
+    "cc = 1.28\n",
+    "sts = [(pc1, ps1)] + [(pc2[y], ps2[y]) for y in range(_A)]\n",
+    "viol = max(float(_np.maximum(a - cc * b, 0).sum() - 0.5 * _np.abs(a - b).sum()) for a, b in sts)\n",
+    "check(\"8e. sum_a (pi - c mu)_+ <= TV(pi, mu) at all 5 states (c=1.28)\", max(viol, 0.0), 0.0)\n",
+    "g_clip = _tok_grad(th_smp, clip=cc)\n",
+    "bound = 0.0\n",
+    "for p, y1, y2, r in _episodes(th_smp):\n",
+    "    q1, q2 = _rho(y1, y2); s1, s2 = _score(th_cur, y1, y2)\n",
+    "    bound += p * abs(r) * (max(q1 - cc, 0) * _np.linalg.norm(s1) + max(q2 - cc, 0) * _np.linalg.norm(s2))\n",
+    "bias_clip = float(_np.linalg.norm(g_clip - g_tok))\n",
+    "print(f\"   clip c={cc}: ||g_clip - g_token-IS|| = {bias_clip:.3e} <= bound {bound:.3e}\")\n",
+    "check(\"8e. clipping bias <= E_mu[sum_t (rho_t - c)_+ |A| ||s_t||]\", max(bias_clip - bound, 0.0), 0.0)\n",
+    "# (f) outcome broadcast vs per-turn credit (on-policy, oracle baselines): both unbiased, different variance\n",
+    "eps_ = _episodes(th_cur)\n",
+    "V = sum(p * r for p, y1, y2, r in eps_)\n",
+    "Q = _np.array([sum(pc2[y1, y2] * _T[y1, y2] for y2 in range(_A)) for y1 in range(_A)])\n",
+    "def _moment(coef1, coef2):\n",
+    "    m1 = 0; m2 = 0\n",
+    "    for p, y1, y2, r in eps_:\n",
+    "        s1, s2 = _score(th_cur, y1, y2); g = coef1(y1, y2, r) * s1 + coef2(y1, y2, r) * s2\n",
+    "        m1 = m1 + p * g; m2 += p * float(g @ g)\n",
+    "    return m1, m2\n",
+    "gb_mean, gb_m2 = _moment(lambda y1, y2, r: r - V, lambda y1, y2, r: r - V)\n",
+    "gt_mean, gt_m2 = _moment(lambda y1, y2, r: Q[y1] - V, lambda y1, y2, r: r - Q[y1])\n",
+    "check(\"8f. broadcast estimator (R-V)(s1+s2) is unbiased for grad J\", float(_np.abs(gb_mean - g_true).max()), 1e-8)\n",
+    "check(\"8f. per-turn estimator (Q-V) s1 + (R-Q) s2 is unbiased for grad J\", float(_np.abs(gt_mean - g_true).max()), 1e-8)\n",
+    "v_b = gb_m2 - float(gb_mean @ gb_mean); v_t = gt_m2 - float(gt_mean @ gt_mean)\n",
+    "print(f\"   total variance (trace): broadcast {v_b:.4f}, per-turn {v_t:.4f}, ratio {v_b/v_t:.3f}\")\n",
+    "# turn-1 decomposition: E[(R-V)^2 ||s1||^2] = E[(Q-V)^2 ||s1||^2] + E[Var(R|y1) ||s1||^2]\n",
+    "lhs = sum(p * (r - V) ** 2 * float(_score(th_cur, y1, y2)[0] @ _score(th_cur, y1, y2)[0]) for p, y1, y2, r in eps_)\n",
+    "t1 = sum(p * (Q[y1] - V) ** 2 * float(_score(th_cur, y1, y2)[0] @ _score(th_cur, y1, y2)[0]) for p, y1, y2, r in eps_)\n",
+    "varR = _np.array([sum(pc2[y1, y2] * (_T[y1, y2] + 0) for y2 in range(_A)) - Q[y1] ** 2 for y1 in range(_A)])   # Var(R|y1) = E[R^2|y1] - Q^2 (R binary: E R^2 = E R)\n",
+    "t2 = sum(p * varR[y1] * float(_score(th_cur, y1, y2)[0] @ _score(th_cur, y1, y2)[0]) for p, y1, y2, r in eps_)\n",
+    "check(\"8f. turn-1 second moment: E[(R-V)^2||s1||^2] = E[(Q-V)^2||s1||^2] + E[Var(R|y1)||s1||^2]\", abs(lhs - t1 - t2), 1e-13)\n",
+    "print(f\"   turn-1: broadcast {lhs:.4f} = {t1:.4f} (credit) + {t2:.4f} (noise from turn 2)\")\n",
+    "check(\"8f. martingale orthogonality: E[s1.s2] = 0 so E||s1+s2||^2 = E||s1||^2 + E||s2||^2\", abs(sum(p * float(_score(th_cur, y1, y2)[0] @ _score(th_cur, y1, y2)[1]) for p, y1, y2, r in eps_)), 1e-15)\n",
+    "# (g) score second moment grows linearly with the number of policy tokens: L-step context-dependent chain, enumerated\n",
+    "Lmax = 5; Th = _rng.normal(0, 0.8, (Lmax, _A, _A)); P0 = _sm(_rng.normal(0, 0.8, _A))\n",
+    "def chain_moments(L):\n",
+    "    tot = 0.0; per = _np.zeros(L); cross = 0.0\n",
+    "    for seq in _np.ndindex(*([_A] * L)):\n",
+    "        pr = P0[seq[0]]; sc = []\n",
+    "        for t_ in range(L):\n",
+    "            prev = 0 if t_ == 0 else seq[t_ - 1]\n",
+    "            pt = P0 if t_ == 0 else _sm(Th[t_, prev])\n",
+    "            if t_ > 0: pr *= pt[seq[t_]]\n",
+    "            v = _np.zeros(L * _A * _A + _A)\n",
+    "            v[:_A] = (_np.eye(_A)[seq[0]] - P0) if t_ == 0 else 0\n",
+    "            if t_ > 0: v[_A + t_ * _A * _A + prev * _A: _A + t_ * _A * _A + (prev + 1) * _A] = _np.eye(_A)[seq[t_]] - pt\n",
+    "            sc.append(v)\n",
+    "        S_ = sum(sc); tot += pr * float(S_ @ S_)\n",
+    "        for t_ in range(L): per[t_] += pr * float(sc[t_] @ sc[t_])\n",
+    "    return tot, per\n",
+    "rows = []\n",
+    "for L in range(1, Lmax + 1):\n",
+    "    tot, per = chain_moments(L); rows.append((L, tot, per.sum()))\n",
+    "    print(f\"   L={L}: E||sum_t s_t||^2 = {tot:.4f}   sum_t E||s_t||^2 = {per.sum():.4f}\")\n",
+    "check(\"8g. E||sum_t s_t||^2 = sum_t E||s_t||^2 for L = 1..5 (score terms are martingale differences)\", max(abs(a - b) for _, a, b in rows), 1e-12)\n",
+    "print(\"done; checks:\", len(CHECKS))"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|8a. E_pi[R (s1+s2)] = finite-difference grad J|2.665e-11|tol=1.0e-08|PASS\n",
+      "CHECK|exact|8b. E_mu[rho1 rho2 R (s1+s2)] = grad J exactly (sampler ratio, full sequence weight)|2.665e-11|tol=1.0e-09|PASS\n",
+      "   token-level sampler ratio, unclipped: ||bias|| = 8.068e-03 (||grad J|| = 1.392e-01)\n",
+      "CHECK|exact|8c. token-level weights are biased for a mismatched sampler (> 1e-4)|0.000e+00|tol=0.0e+00|PASS\n",
+      "CHECK|exact|8c. ... and exact when the sampler equals the policy|2.665e-11|tol=1.0e-08|PASS\n",
+      "CHECK|exact|8d. constant baseline b adds exactly 0 to the token-level sampler-ratio gradient: max |E_mu[sum_t rho_t b s_t]| over 20 coordinates|2.776e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|8e. sum_a (pi - c mu)_+ <= TV(pi, mu) at all 5 states (c=1.28)|0.000e+00|tol=0.0e+00|PASS\n",
+      "   clip c=1.28: ||g_clip - g_token-IS|| = 2.023e-02 <= bound 3.210e-02\n",
+      "CHECK|exact|8e. clipping bias <= E_mu[sum_t (rho_t - c)_+ |A| ||s_t||]|0.000e+00|tol=0.0e+00|PASS\n",
+      "CHECK|exact|8f. broadcast estimator (R-V)(s1+s2) is unbiased for grad J|2.665e-11|tol=1.0e-08|PASS\n",
+      "CHECK|exact|8f. per-turn estimator (Q-V) s1 + (R-Q) s2 is unbiased for grad J|2.665e-11|tol=1.0e-08|PASS\n",
+      "   total variance (trace): broadcast 0.3020, per-turn 0.1473, ratio 2.050\n",
+      "CHECK|exact|8f. turn-1 second moment: E[(R-V)^2||s1||^2] = E[(Q-V)^2||s1||^2] + E[Var(R|y1)||s1||^2]|5.551e-17|tol=1.0e-13|PASS\n",
+      "   turn-1: broadcast 0.1641 = 0.0381 (credit) + 0.1260 (noise from turn 2)\n",
+      "CHECK|exact|8f. martingale orthogonality: E[s1.s2] = 0 so E||s1+s2||^2 = E||s1||^2 + E||s2||^2|0.000e+00|tol=1.0e-15|PASS\n",
+      "   L=1: E||sum_t s_t||^2 = 0.7188   sum_t E||s_t||^2 = 0.7188\n",
+      "   L=2: E||sum_t s_t||^2 = 1.4115   sum_t E||s_t||^2 = 1.4115\n",
+      "   L=3: E||sum_t s_t||^2 = 2.0546   sum_t E||s_t||^2 = 2.0546\n",
+      "   L=4: E||sum_t s_t||^2 = 2.6447   sum_t E||s_t||^2 = 2.6447\n",
+      "   L=5: E||sum_t s_t||^2 = 3.2724   sum_t E||s_t||^2 = 3.2724\n",
+      "CHECK|exact|8g. E||sum_t s_t||^2 = sum_t E||s_t||^2 for L = 1..5 (score terms are martingale differences)|5.329e-15|tol=1.0e-12|PASS\n",
+      "done; checks: 28\n"
+     ]
+    }
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

code/notebooks/0108-training-inference-mismatch.ipynb

@@ -0,0 +1,585 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 0108 - Training-inference mismatch: sampler vs trainer log-probs, truncated IS, why an FP32 LM head helps\n",
+    "Companion to `notes/0108 - Training-Inference Mismatch.pdf`.  fp64 reference (\"trainer\" pi); the \"sampler\" mu uses logits rounded to bf16/fp16/fp32."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "source": [
+    "import sys, math\n",
+    "sys.path.insert(0, \"..\")\n",
+    "from _m1_lib import *\n",
+    "import numpy as np"
+   ],
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 0. Floating-point facts: unit roundoff u = 2^-(p) with p significant bits;  |round(x) - x| <= u |x|\n",
+    "bf16: 8 significant bits (u=2^-8), fp16: 11 (u=2^-11), fp32: 24 (u=2^-24)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "source": [
+    "gen = torch.Generator().manual_seed(0)\n",
+    "x = torch.exp(torch.rand(2_000_000, generator=gen) * math.log(1e3 / 1e-3) + math.log(1e-3))    # log-uniform in [1e-3, 1e3]\n",
+    "for name, dt, u in ((\"bf16\", torch.bfloat16, 2.0 ** -8), (\"fp16\", torch.float16, 2.0 ** -11), (\"fp32\", torch.float32, 2.0 ** -24)):\n",
+    "    rel = ((x.to(dt).to(torch.float64) - x).abs() / x)\n",
+    "    print(f\"   {name}: max relative rounding error {float(rel.max()):.3e}   u = {u:.3e}\")\n",
+    "    check(f\"{name}: max rel. rounding error <= u (=2^-{int(-math.log2(u))})\", max(0.0, float(rel.max()) - u), 0.0)\n",
+    "    assert float(rel.max()) > 0.5 * u"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "   bf16: max relative rounding error 3.891e-03   u = 3.906e-03\n",
+      "CHECK|exact|bf16: max rel. rounding error <= u (=2^-8)|0.000e+00|tol=0.0e+00|PASS\n",
+      "   fp16: max relative rounding error 4.880e-04   u = 4.883e-04\n",
+      "CHECK|exact|fp16: max rel. rounding error <= u (=2^-11)|0.000e+00|tol=0.0e+00|PASS\n",
+      "   fp32: max relative rounding error 5.958e-08   u = 5.960e-08\n",
+      "CHECK|exact|fp32: max rel. rounding error <= u (=2^-24)|0.000e+00|tol=0.0e+00|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. The tiny LM with LARGE logits (like an LM head with high-magnitude activations); sampler mu = same network with bf16-rounded logits"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "source": [
+    "lm = TinyLM(V=4, L=4, seed=0, scale=4.0)\n",
+    "theta = lm.theta0.clone()\n",
+    "R = default_reward(lm)\n",
+    "def seq_logprob_rounded(theta, dtype):\n",
+    "    V, L = lm.V, lm.L\n",
+    "    tok, pos = lm._logits_tables(theta)\n",
+    "    prev = torch.cat([torch.full((lm.N, 1), V, dtype=torch.long), lm.tokens[:, :-1]], 1)\n",
+    "    logits = tok[prev] + pos[None, :, :]\n",
+    "    if dtype is not None: logits = logits.to(dtype).to(torch.float64)\n",
+    "    lp = torch.log_softmax(logits, -1).gather(-1, lm.tokens[..., None]).squeeze(-1) * lm.mask\n",
+    "    return lp.sum(1)\n",
+    "pi = seq_logprob_rounded(theta, None).exp()\n",
+    "mu = seq_logprob_rounded(theta, torch.bfloat16).exp()\n",
+    "check(\"trainer pi is normalised\", abs(float(pi.sum()) - 1), 1e-14)\n",
+    "check(\"sampler mu is normalised (bf16 logits then softmax)\", abs(float(mu.sum()) - 1), 1e-14)\n",
+    "b = float(pi @ R); A = R - b\n",
+    "S = score_matrix(lm, theta)                                  # score of the TRAINER policy\n",
+    "g_true = (pi[:, None] * A[:, None] * S).sum(0)               # grad J_pi\n",
+    "w = pi / mu\n",
+    "TVs = 0.5 * float((pi - mu).abs().sum())\n",
+    "print(f\"   bf16 logits: TV(pi, mu) = {TVs:.5f}   chi2(pi||mu) = {float((mu*(w-1)**2).sum()):.3e}   max w = {float(w.max()):.4f}\")\n",
+    "for nm, dt in ((\"fp16\", torch.float16), (\"fp32\", torch.float32)):\n",
+    "    m_ = seq_logprob_rounded(theta, dt).exp()\n",
+    "    print(f\"   {nm} logits: TV(pi, mu) = {0.5*float((pi-m_).abs().sum()):.3e}   chi2 = {float((m_*(pi/m_-1)**2).sum()):.3e}\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|trainer pi is normalised|0.000e+00|tol=1.0e-14|PASS\n",
+      "CHECK|exact|sampler mu is normalised (bf16 logits then softmax)|0.000e+00|tol=1.0e-14|PASS\n",
+      "   bf16 logits: TV(pi, mu) = 0.00045   chi2(pi||mu) = 5.187e-06   max w = 1.0491\n",
+      "   fp16 logits: TV(pi, mu) = 1.195e-04   chi2 = 2.142e-07\n",
+      "   fp32 logits: TV(pi, mu) = 4.361e-09   chi2 = 1.076e-15\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Ignoring the mismatch is biased; IS repairs it exactly  (arXiv:2510.26788 Eqs. 3, 5)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "source": [
+    "g_biased = (mu[:, None] * A[:, None] * S).sum(0)             # samples from mu, trainer's score, no correction\n",
+    "bias_formula = ((mu - pi)[:, None] * A[:, None] * S).sum(0)  # = E_pi[(mu/pi - 1) A s]\n",
+    "check(\"bias of the uncorrected estimator == sum_o (mu-pi) A s\", maxabs(g_biased - g_true, bias_formula), 1e-14)\n",
+    "g_is = (mu[:, None] * w[:, None] * A[:, None] * S).sum(0)\n",
+    "check(\"sequence-level IS correction E_mu[(pi/mu) A s] == grad J_pi exactly\", maxabs(g_is, g_true), 1e-14)\n",
+    "print(f\"   ||bias|| / ||grad J|| = {float((g_biased-g_true).norm()/g_true.norm()):.4f}   (uncorrected)\")\n",
+    "assert float((g_biased - g_true).norm()) > 1e-4"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|bias of the uncorrected estimator == sum_o (mu-pi) A s|1.930e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|sequence-level IS correction E_mu[(pi/mu) A s] == grad J_pi exactly|6.939e-18|tol=1.0e-14|PASS\n",
+      "   ||bias|| / ||grad J|| = 0.0011   (uncorrected)\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Truncated IS (TIS, Eq. 6) and masked IS (MIS, Eq. 7): exact bias and bounds\n",
+    "  TIS: E_mu[min(w,C) A s];   bias = - E_mu[(w-C)_+ A s]       ||.|| <= M_A M_s sum_o (pi - C mu)_+  <= M_A M_s min(TV, P_pi(w>C))\n",
+    "  MIS: E_mu[1(w<=C) w A s];  bias = - E_pi[1(w>C) A s]         ||.|| <= M_A M_s P_pi(w>C)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "source": [
+    "MA, Ms = float(A.abs().max()), float(S.norm(dim=1).max())\n",
+    "print(\"   C       ||TIS bias||   bound1(sum mu (w-C)+ |A| ||s||)   bound2(MA Ms sum(pi-C mu)+)   MA Ms TV    ||MIS bias||   MA Ms P_pi(w>C)\")\n",
+    "for C in (1.0, 1.02, 1.05, 1.2, 2.0):\n",
+    "    wt = torch.clamp(w, max=C)\n",
+    "    g_tis = (mu[:, None] * wt[:, None] * A[:, None] * S).sum(0)\n",
+    "    bias_tis = g_tis - g_true\n",
+    "    formula = -((mu * (w - C).clamp(min=0))[:, None] * A[:, None] * S).sum(0)\n",
+    "    check(f\"TIS bias == -E_mu[(w-C)_+ A s]  (C={C})\", maxabs(bias_tis, formula), 1e-14)\n",
+    "    b1 = float((mu * (w - C).clamp(min=0) * A.abs() * S.norm(dim=1)).sum())\n",
+    "    ex = float((pi - C * mu).clamp(min=0).sum())\n",
+    "    b2 = MA * Ms * ex\n",
+    "    assert float(bias_tis.norm()) <= b1 + 1e-15 and b1 <= b2 + 1e-15 and ex <= TVs + 1e-15\n",
+    "    g_mis = (mu[:, None] * (w * (w <= C))[:, None] * A[:, None] * S).sum(0)\n",
+    "    bias_mis = g_mis - g_true\n",
+    "    mis_formula = -((pi * (w > C))[:, None] * A[:, None] * S).sum(0)\n",
+    "    check(f\"MIS bias == -E_pi[1(w>C) A s]  (C={C})\", maxabs(bias_mis, mis_formula), 1e-14)\n",
+    "    Pw = float((pi * (w > C)).sum())\n",
+    "    assert float(bias_mis.norm()) <= MA * Ms * Pw + 1e-15\n",
+    "    print(f\"   {C:4.2f}   {float(bias_tis.norm()):12.4e}   {b1:12.4e}                  {b2:12.4e}              {MA*Ms*TVs:10.4e}  {float(bias_mis.norm()):12.4e}   {MA*Ms*Pw:12.4e}\")\n",
+    "check(\"TIS: 0 <= bias <= bound1 <= bound2 <= M_A M_s TV(pi,mu) for all C\", 0.0, 0.0)\n",
+    "check(\"TIS at C=inf is exact (no truncation) -> bias = 0\", float((((mu[:, None] * torch.clamp(w, max=1e12)[:, None] * A[:, None] * S).sum(0)) - g_true).norm()), 1e-13)\n",
+    "# TV of sequence distributions <= sum over positions of E[TV_t]   (hybrid argument)\n",
+    "def tv_chain(theta, dtype):\n",
+    "    tot = 0.0\n",
+    "    V, L = lm.V, lm.L\n",
+    "    tok, pos = lm._logits_tables(theta)\n",
+    "    def cond(prefix, rounder):\n",
+    "        prev = V if len(prefix) == 0 else prefix[-1]\n",
+    "        z = tok[prev] + pos[len(prefix)]\n",
+    "        if rounder is not None: z = z.to(rounder).to(torch.float64)\n",
+    "        return torch.softmax(z, 0)\n",
+    "    def rec(prefix, prob):\n",
+    "        nonlocal tot\n",
+    "        if len(prefix) >= L: return\n",
+    "        cp, cm = cond(prefix, None), cond(prefix, dtype)\n",
+    "        tot += prob * 0.5 * float((cp - cm).abs().sum())\n",
+    "        for a in range(V - 1): rec(prefix + (a,), prob * float(cp[a]))\n",
+    "    rec((), 1.0)\n",
+    "    return tot\n",
+    "tvc = tv_chain(theta, torch.bfloat16)\n",
+    "print(f\"   TV(sequence) = {TVs:.4f} <= sum_t E_pi[TV_t] = {tvc:.4f}\")\n",
+    "check(\"sequence TV <= sum_t E[TV_t]  (max violation, must be <= 0)\", max(0.0, TVs - tvc), 0.0)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "   C       ||TIS bias||   bound1(sum mu (w-C)+ |A| ||s||)   bound2(MA Ms sum(pi-C mu)+)   MA Ms TV    ||MIS bias||   MA Ms P_pi(w>C)\n",
+      "CHECK|exact|TIS bias == -E_mu[(w-C)_+ A s]  (C=1.0)|2.557e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|MIS bias == -E_pi[1(w>C) A s]  (C=1.0)|6.939e-18|tol=1.0e-14|PASS\n",
+      "   1.00     1.1248e-04     1.6072e-04                    2.2052e-03              2.2052e-03    1.2226e-01     4.0596e+00\n",
+      "CHECK|exact|TIS bias == -E_mu[(w-C)_+ A s]  (C=1.02)|2.005e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|MIS bias == -E_pi[1(w>C) A s]  (C=1.02)|1.562e-17|tol=1.0e-14|PASS\n",
+      "   1.02     6.4034e-08     6.6273e-08                    1.6801e-06              2.2052e-03    4.3457e-05     1.1315e-03\n",
+      "CHECK|exact|TIS bias == -E_mu[(w-C)_+ A s]  (C=1.05)|6.939e-18|tol=1.0e-14|PASS\n",
+      "CHECK|exact|MIS bias == -E_pi[1(w>C) A s]  (C=1.05)|6.939e-18|tol=1.0e-14|PASS\n",
+      "   1.05     1.7358e-17     0.0000e+00                    0.0000e+00              2.2052e-03    1.7358e-17     0.0000e+00\n",
+      "CHECK|exact|TIS bias == -E_mu[(w-C)_+ A s]  (C=1.2)|6.939e-18|tol=1.0e-14|PASS\n",
+      "CHECK|exact|MIS bias == -E_pi[1(w>C) A s]  (C=1.2)|6.939e-18|tol=1.0e-14|PASS\n",
+      "   1.20     1.7358e-17     0.0000e+00                    0.0000e+00              2.2052e-03    1.7358e-17     0.0000e+00\n",
+      "CHECK|exact|TIS bias == -E_mu[(w-C)_+ A s]  (C=2.0)|6.939e-18|tol=1.0e-14|PASS\n",
+      "CHECK|exact|MIS bias == -E_pi[1(w>C) A s]  (C=2.0)|6.939e-18|tol=1.0e-14|PASS\n",
+      "   2.00     1.7358e-17     0.0000e+00                    0.0000e+00              2.2052e-03    1.7358e-17     0.0000e+00\n",
+      "CHECK|exact|TIS: 0 <= bias <= bound1 <= bound2 <= M_A M_s TV(pi,mu) for all C|0.000e+00|tol=0.0e+00|PASS\n",
+      "CHECK|exact|TIS at C=inf is exact (no truncation) -> bias = 0|1.736e-17|tol=1.0e-13|PASS\n",
+      "   TV(sequence) = 0.0004 <= sum_t E_pi[TV_t] = 0.0007\n",
+      "CHECK|exact|sequence TV <= sum_t E[TV_t]  (max violation, must be <= 0)|0.000e+00|tol=0.0e+00|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. From logit rounding error to log-prob error:  d log pi_a = e_a - sum_b pi_b e_b  (first order);  Var = s^2 (1 - 2 pi_a + sum pi_b^2) for iid logit errors"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "source": [
+    "gen = torch.Generator().manual_seed(3)\n",
+    "z = torch.randn(7, generator=gen) * 5.0\n",
+    "pi_z = torch.softmax(z, 0)\n",
+    "eps = torch.randn(7, generator=gen) * 1e-6\n",
+    "dlp_exact = torch.log_softmax(z + eps, 0) - torch.log_softmax(z, 0)\n",
+    "dlp_lin = eps - (pi_z * eps).sum()\n",
+    "check(\"first-order log-prob error  e_a - sum_b pi_b e_b  (relative error of linearisation, e~1e-6)\", float((dlp_exact - dlp_lin).abs().max() / dlp_lin.abs().max()), 1e-5)\n",
+    "s_ = 0.01\n",
+    "E = torch.randn(400000, 7, generator=gen) * s_\n",
+    "dlp = torch.log_softmax(z + E, -1) - torch.log_softmax(z, 0)\n",
+    "var_mc = dlp.var(0); var_f = s_ ** 2 * (1 - 2 * pi_z + (pi_z ** 2).sum())\n",
+    "zs = float(((var_mc - var_f).abs() / (var_f * math.sqrt(2 / 400000))).max())\n",
+    "check(\"Var(d log pi_a) = s^2 (1 - 2 pi_a + sum pi^2): max z-score over 7 tokens (MC, 400k draws)\", zs, 4.0, kind=\"mc\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|first-order log-prob error  e_a - sum_b pi_b e_b  (relative error of linearisation, e~1e-6)|5.237e-08|tol=1.0e-05|PASS\n",
+      "CHECK|mc|Var(d log pi_a) = s^2 (1 - 2 pi_a + sum pi^2): max z-score over 7 tokens (MC, 400k draws)|1.196e+00|tol=4.0e+00|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4b. Where the logit error comes from: an LM-head matmul with bf16 inputs vs an fp32 head\n",
+    "z_b = sum_i h_i W_bi.  Rounding h and W to bf16 perturbs each product by a relative error <= u, so the logit error scales with the\n",
+    "magnitude of the terms (not just of z).  Outlier activations (a few huge hidden dims) enlarge that magnitude."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "source": [
+    "gen = torch.Generator().manual_seed(4)\n",
+    "d, Vv, Tn = 1024, 64, 512\n",
+    "def head_errors(outlier):\n",
+    "    h = torch.randn(Tn, d, generator=gen)\n",
+    "    h[:, :8] *= outlier                                         # 8 high-magnitude hidden dimensions\n",
+    "    W = torch.randn(Vv, d, generator=gen) / math.sqrt(d)\n",
+    "    z = h @ W.T                                                 # fp64 reference\n",
+    "    zb = h.to(torch.bfloat16).to(torch.float64) @ W.to(torch.bfloat16).to(torch.float64).T     # bf16-rounded inputs, exact accumulation\n",
+    "    z32 = h.to(torch.float32).to(torch.float64) @ W.to(torch.float32).to(torch.float64).T      # fp32-rounded inputs\n",
+    "    return float(z.std()), float((zb - z).std()), float((z32 - z).std())\n",
+    "print(\"   outlier x    std(z)    std(err bf16)   std(err fp32)    err_bf16/std(z)   /u_bf16     ratio bf16/fp32\")\n",
+    "for out in (1.0, 10.0, 100.0):\n",
+    "    sz_, eb, e32 = head_errors(out)\n",
+    "    print(f\"   {out:7.0f}   {sz_:8.2f}   {eb:12.3e}   {e32:12.3e}      {eb/sz_:10.3e}    {eb/sz_/2**-8:6.3f}     {eb/e32:9.1f}\")\n",
+    "sz_, eb, e32 = head_errors(100.0)\n",
+    "check(\"bf16 head: logit error / std(z) between 0.05u and 1u  (u=2^-8)\", 0.0 if 0.05 * 2**-8 < eb / sz_ < 2**-8 else 1.0, 0.0)\n",
+    "check(\"fp32 head cuts the logit error by ~ u_bf16/u_fp32 = 65536 (log10 ratio within 1 decade)\", abs(math.log10(eb / e32) - math.log10(2**16)), 1.0)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "   outlier x    std(z)    std(err bf16)   std(err fp32)    err_bf16/std(z)   /u_bf16     ratio bf16/fp32\n",
+      "         1       1.01      2.366e-03      3.575e-08       2.349e-03     0.601       66166.0\n",
+      "        10       1.34      3.054e-03      4.786e-08       2.284e-03     0.585       63825.0\n",
+      "       100       8.81      2.025e-02      3.259e-07       2.298e-03     0.588       62139.1\n",
+      "CHECK|exact|bf16 head: logit error / std(z) between 0.05u and 1u  (u=2^-8)|0.000e+00|tol=0.0e+00|PASS\n",
+      "CHECK|exact|fp32 head cuts the logit error by ~ u_bf16/u_fp32 = 65536 (log10 ratio within 1 decade)|2.678e-02|tol=1.0e+00|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Long horizons: positions are independent contexts, so everything is EXACT:\n",
+    "  E_mu[log w_T] = -sum_t KL(mu_t||pi_t),  Var(log w_T) = sum_t Var_t,  E_mu[w_T^2] = prod_t (1 + chi2_t).   (V=64, logits ~ N(0, sigma_z^2))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "source": [
+    "def per_position_stats(sigma_z, dtype, T=4096, V=64, seed=0):\n",
+    "    g = torch.Generator().manual_seed(seed)\n",
+    "    z = torch.randn(T, V, generator=g) * sigma_z\n",
+    "    pi = torch.softmax(z, -1)\n",
+    "    zr = z if dtype is None else z.to(dtype).to(torch.float64)\n",
+    "    mu = torch.softmax(zr, -1)\n",
+    "    lw = pi.log() - mu.log()                                  # log(pi/mu) per token option\n",
+    "    Elw = (mu * lw).sum(-1)                                   # E_mu[log w_t] = -KL(mu||pi)\n",
+    "    Vlw = (mu * lw ** 2).sum(-1) - Elw ** 2\n",
+    "    chi2 = (mu * (pi / mu - 1) ** 2).sum(-1)\n",
+    "    return Elw, Vlw, chi2, pi, mu\n",
+    "print(\"   sigma_z  head      mean KL/token    mean chi2/token     T=100: ESS frac     T=1000      T=4000   (ESS frac = 1/prod(1+chi2))\")\n",
+    "rows = {}\n",
+    "for sz in (2.0, 5.0, 10.0, 20.0):\n",
+    "    for name, dt in ((\"bf16\", torch.bfloat16), (\"fp16\", torch.float16), (\"fp32\", torch.float32)):\n",
+    "        Elw, Vlw, chi2, pi_t, mu_t = per_position_stats(sz, dt)\n",
+    "        lESS = lambda T: float(-torch.log1p(chi2[:T]).sum())\n",
+    "        rows[(sz, name)] = (float((-Elw).mean()), float(chi2.mean()), lESS(100), lESS(1000), lESS(4000))\n",
+    "        print(f\"   {sz:5.1f}   {name}   {float((-Elw).mean()):14.3e}   {float(chi2.mean()):14.3e}      {math.exp(lESS(100)):10.3e}   {math.exp(lESS(1000)):10.3e}  {math.exp(lESS(4000)):10.3e}\")\n",
+    "# why chi2 does not grow like sigma_z^2: to first order chi2_t = Var_{a~pi}(e_a) where e = round(z) - z (the *centred* logit error)\n",
+    "def first_order_check(sz, dtype, u, T=2048, V=64):\n",
+    "    g_ = torch.Generator().manual_seed(0)\n",
+    "    z_ = torch.randn(T, V, generator=g_) * sz\n",
+    "    e_ = z_.to(dtype).to(torch.float64) - z_\n",
+    "    pi_ = torch.softmax(z_, -1); mu_ = torch.softmax(z_ + e_, -1)\n",
+    "    chi2_ = (mu_ * (pi_ / mu_ - 1) ** 2).sum(-1)\n",
+    "    var_e = (pi_ * e_ ** 2).sum(-1) - ((pi_ * e_).sum(-1)) ** 2\n",
+    "    bound = u ** 2 * (pi_ * z_ ** 2).sum(-1)\n",
+    "    return float(chi2_.sum() / var_e.sum()), float((var_e / bound).max()), float(chi2_.mean())\n",
+    "for sz in (2.0, 10.0, 20.0):\n",
+    "    ratio, bnd, c2 = first_order_check(sz, torch.bfloat16, 2.0 ** -8)\n",
+    "    print(f\"   sigma_z={sz:5.1f}: chi2/Var_pi(e) = {ratio:.4f}   max Var_pi(e)/(u^2 E_pi z^2) = {bnd:.3f}\")\n",
+    "    check(f\"first-order law chi2_t = Var_pi(logit rounding error)  (sigma_z={sz}, mean ratio - 1)\", abs(ratio - 1), 0.1)\n",
+    "    check(f\"bound Var_pi(e) <= u^2 E_pi[z^2]  (sigma_z={sz}, max ratio - 1 must be <= 0)\", max(0.0, bnd - 1), 0.0)\n",
+    "szs = np.array([2.0, 5.0, 10.0, 20.0]); chis = np.array([rows[(s, 'bf16')][1] for s in szs])\n",
+    "slope = np.polyfit(np.log(szs), np.log(chis), 1)[0]\n",
+    "print(f\"   bf16 log-log slope of chi2/token vs logit scale: {slope:.3f}  (between 1 and 2: the error grows like |z| but the centred error Var_pi(e) shrinks as pi becomes peaked)\")\n",
+    "assert 1.0 < slope < 2.0\n",
+    "ratio_fp32 = rows[(10.0, 'fp32')][1] / rows[(10.0, 'bf16')][1]\n",
+    "print(f\"   fp32/bf16 chi2 ratio at sigma_z=10: {ratio_fp32:.2e}   (theory (u32/u16)^2 = {(2.0**-24/2.0**-8)**2:.2e})\")\n",
+    "check(\"fp32 head reduces per-token chi^2 by ~ (u_fp32/u_bf16)^2  (|log10 ratio - log10 theory|)\", abs(math.log10(ratio_fp32) - math.log10((2.0 ** -24 / 2.0 ** -8) ** 2)), 1.0)\n",
+    "# MC check of one exact statistic: Var(log w_T) over sampled sequences for T=64\n",
+    "Elw, Vlw, chi2, pi_t, mu_t = per_position_stats(10.0, torch.bfloat16, T=64)\n",
+    "g = torch.Generator().manual_seed(9); n = 20000\n",
+    "tokens = torch.stack([torch.multinomial(mu_t[t], n, replacement=True, generator=g) for t in range(64)], 1)       # [n, T]\n",
+    "lw_seq = (pi_t.log().gather(1, tokens.T).T - mu_t.log().gather(1, tokens.T).T).sum(1)\n",
+    "exact_mean, exact_var = float(Elw.sum()), float(Vlw.sum())\n",
+    "se_mean = math.sqrt(exact_var / n)\n",
+    "check(\"MC: E_mu[log w_64] == -sum KL  (z-score)\", abs(float(lw_seq.mean()) - exact_mean) / se_mean, 4.0, kind=\"mc\")\n",
+    "check(\"MC: Var(log w_64) == sum_t Var_t  (z-score; kurtosis ignored -> loose 6)\", abs(float(lw_seq.var()) - exact_var) / (exact_var * math.sqrt(2 / n)), 6.0, kind=\"mc\")\n",
+    "check(\"E_mu[w_T] = 1 exactly per position (sum_a pi_a = 1)\", float((mu_t * (pi_t / mu_t)).sum(-1).sub(1).abs().max()), 1e-13)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "   sigma_z  head      mean KL/token    mean chi2/token     T=100: ESS frac     T=1000      T=4000   (ESS frac = 1/prod(1+chi2))\n",
+      "     2.0   bf16        1.420e-05        2.841e-05       9.967e-01    9.716e-01   8.925e-01\n",
+      "     2.0   fp16        2.179e-07        4.358e-07       1.000e+00    9.996e-01   9.983e-01\n",
+      "     2.0   fp32        3.269e-15        6.533e-15       1.000e+00    1.000e+00   1.000e+00\n",
+      "     5.0   bf16        6.224e-05        1.245e-04       9.895e-01    8.878e-01   6.068e-01\n",
+      "     5.0   fp16        9.391e-07        1.878e-06       9.998e-01    9.980e-01   9.925e-01\n",
+      "     5.0   fp32        1.424e-14        2.849e-14       1.000e+00    1.000e+00   1.000e+00\n",
+      "    10.0   bf16        1.464e-04        2.929e-04       9.774e-01    7.601e-01   3.091e-01\n",
+      "    10.0   fp16        2.135e-06        4.270e-06       9.996e-01    9.956e-01   9.831e-01\n",
+      "    10.0   fp32        3.233e-14        6.466e-14       1.000e+00    1.000e+00   1.000e+00\n",
+      "    20.0   bf16        3.125e-04        6.249e-04       9.580e-01    5.596e-01   8.265e-02\n",
+      "    20.0   fp16        4.395e-06        8.792e-06       9.991e-01    9.914e-01   9.656e-01\n",
+      "    20.0   fp32        6.719e-14        1.344e-13       1.000e+00    1.000e+00   1.000e+00\n",
+      "   sigma_z=  2.0: chi2/Var_pi(e) = 1.0000   max Var_pi(e)/(u^2 E_pi z^2) = 0.590\n",
+      "CHECK|exact|first-order law chi2_t = Var_pi(logit rounding error)  (sigma_z=2.0, mean ratio - 1)|1.799e-05|tol=1.0e-01|PASS\n",
+      "CHECK|exact|bound Var_pi(e) <= u^2 E_pi[z^2]  (sigma_z=2.0, max ratio - 1 must be <= 0)|0.000e+00|tol=0.0e+00|PASS\n",
+      "   sigma_z= 10.0: chi2/Var_pi(e) = 1.0005   max Var_pi(e)/(u^2 E_pi z^2) = 0.709\n",
+      "CHECK|exact|first-order law chi2_t = Var_pi(logit rounding error)  (sigma_z=10.0, mean ratio - 1)|5.085e-04|tol=1.0e-01|PASS\n",
+      "CHECK|exact|bound Var_pi(e) <= u^2 E_pi[z^2]  (sigma_z=10.0, max ratio - 1 must be <= 0)|0.000e+00|tol=0.0e+00|PASS\n",
+      "   sigma_z= 20.0: chi2/Var_pi(e) = 1.0021   max Var_pi(e)/(u^2 E_pi z^2) = 0.644\n",
+      "CHECK|exact|first-order law chi2_t = Var_pi(logit rounding error)  (sigma_z=20.0, mean ratio - 1)|2.100e-03|tol=1.0e-01|PASS\n",
+      "CHECK|exact|bound Var_pi(e) <= u^2 E_pi[z^2]  (sigma_z=20.0, max ratio - 1 must be <= 0)|0.000e+00|tol=0.0e+00|PASS\n",
+      "   bf16 log-log slope of chi2/token vs logit scale: 1.342  (between 1 and 2: the error grows like |z| but the centred error Var_pi(e) shrinks as pi becomes peaked)\n",
+      "   fp32/bf16 chi2 ratio at sigma_z=10: 2.21e-10   (theory (u32/u16)^2 = 2.33e-10)\n",
+      "CHECK|exact|fp32 head reduces per-token chi^2 by ~ (u_fp32/u_bf16)^2  (|log10 ratio - log10 theory|)|2.316e-02|tol=1.0e+00|PASS\n",
+      "CHECK|mc|MC: E_mu[log w_64] == -sum KL  (z-score)|4.121e-01|tol=4.0e+00|PASS\n",
+      "CHECK|mc|MC: Var(log w_64) == sum_t Var_t  (z-score; kurtosis ignored -> loose 6)|5.377e-01|tol=6.0e+00|PASS\n",
+      "CHECK|exact|E_mu[w_T] = 1 exactly per position (sum_a pi_a = 1)|6.661e-16|tol=1.0e-13|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5b. Caveat: E_mu[w_T^2] = prod_t (1 + chi2_t) needs history-INDEPENDENT chi2_t.  With state-dependent chi2 the exact statements are\n",
+    "  (i) the recursion  E[w_T^2] = E[w_{T-1}^2 (1 + chi2_T(s_T))],   (ii) the martingale identity  E[ w_T^2 / prod_t (1 + chi2_t(s_t)) ] = 1,\n",
+    "and E[prod_t (1 + chi2_t(s_t))] is NOT E[w_T^2].  Exact enumeration on a tiny autoregressive model (sampler mu = one weight vector, trainer pi = a perturbed one)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "source": [
+    "lmx = TinyLM(V=3, L=8, seed=4, scale=1.0)\n",
+    "tokx = lmx.theta0[:lmx.n_tok].reshape(lmx.V + 1, lmx.V); tokx[:, lmx.EOS] -= 2.0\n",
+    "th_mu = lmx.theta0.clone(); th_pi = perturb(th_mu, 0.8, seed=9)\n",
+    "lmu, lpi = lmx.token_logprobs(th_mu), lmx.token_logprobs(th_pi)\n",
+    "mu_x = lmu.sum(1).exp()\n",
+    "chi_x = torch.zeros(lmx.N, lmx.L)                                   # chi2_t(s_t) = chi2(pi(.|s_t) || mu(.|s_t)) at the realised prefix\n",
+    "for i, sq in enumerate(lmx.seqs):\n",
+    "    for t in range(len(sq)):\n",
+    "        cm, cp = lmx.cond_probs(th_mu, sq[:t]), lmx.cond_probs(th_pi, sq[:t])\n",
+    "        chi_x[i, t] = ((cp - cm) ** 2 / cm).sum()\n",
+    "w2_x = (lpi - lmu).sum(1).exp() ** 2\n",
+    "prod_x = (1 + chi_x).prod(1)                                        # masked positions have chi2 = 0 -> factor 1\n",
+    "Ew2 = float((mu_x * w2_x).sum()); Eprod = float((mu_x * prod_x).sum())\n",
+    "w2_prev = torch.cat([torch.zeros(lmx.N, 1), (lpi - lmu).cumsum(1)[:, :-1]], 1).exp() ** 2        # w_{t-1}^2 (w_0 = 1)\n",
+    "rec = float((mu_x * w2_prev[:, -1] * (1 + chi_x[:, -1])).sum())\n",
+    "print(f\"   N={lmx.N}:  E_mu[w^2] = {Ew2:.6f};  recursion E[w_(T-1)^2 (1+chi2_T)] = {rec:.6f};  E_mu[prod_t(1+chi2_t)] = {Eprod:.6f};  E_mu[w^2/prod] = {float((mu_x*w2_x/prod_x).sum()):.6f}\")\n",
+    "check(\"recursion E[w_T^2] = E[w_(T-1)^2 (1+chi2_T(s_T))]  (state-dependent chi2)\", abs(Ew2 - rec), 1e-14)\n",
+    "check(\"martingale identity E[w_T^2 / prod_t (1+chi2_t(s_t))] = 1\", abs(float((mu_x * w2_x / prod_x).sum()) - 1.0), 1e-13)\n",
+    "assert abs(Ew2 - Eprod) > 1e-3          # the naive product formula is genuinely wrong when chi2 depends on the history\n",
+    "check(\"E[prod(1+chi2_t)] != E[w_T^2] on a state-dependent model (relative gap, must exceed 1e-3)\", max(0.0, 1e-3 - abs(Ew2 - Eprod) / Ew2), 0.0)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "   N=511:  E_mu[w^2] = 1.151486;  recursion E[w_(T-1)^2 (1+chi2_T)] = 1.151486;  E_mu[prod_t(1+chi2_t)] = 1.141717;  E_mu[w^2/prod] = 1.000000\n",
+      "CHECK|exact|recursion E[w_T^2] = E[w_(T-1)^2 (1+chi2_T(s_T))]  (state-dependent chi2)|6.661e-16|tol=1.0e-14|PASS\n",
+      "CHECK|exact|martingale identity E[w_T^2 / prod_t (1+chi2_t(s_t))] = 1|4.441e-16|tol=1.0e-13|PASS\n",
+      "CHECK|exact|E[prod(1+chi2_t)] != E[w_T^2] on a state-dependent model (relative gap, must exceed 1e-3)|0.000e+00|tol=0.0e+00|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Consequence for PPO/CISPO ratios and for a \"probability correlation\" diagnostic\n",
+    "If the *sampler's* log-probs are used as pi_old, then at theta = theta_old the ratio is r = pi/mu != 1: tokens are clipped before any update."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "source": [
+    "for name, dt in ((\"bf16\", torch.bfloat16), (\"fp32\", torch.float32)):\n",
+    "    Elw, Vlw, chi2, pi_t, mu_t = per_position_stats(10.0, dt, T=2048)\n",
+    "    gen = torch.Generator().manual_seed(1)\n",
+    "    tok_ = torch.stack([torch.multinomial(mu_t[t], 64, replacement=True, generator=gen) for t in range(2048)], 0)   # [T, 64]\n",
+    "    r = (pi_t.gather(1, tok_) / mu_t.gather(1, tok_))\n",
+    "    out = float(((r - 1).abs() > 0.2).double().mean())\n",
+    "    p_tr, p_inf = pi_t.gather(1, tok_).reshape(-1), mu_t.gather(1, tok_).reshape(-1)\n",
+    "    corr = float(torch.corrcoef(torch.stack([p_tr, p_inf]))[0, 1])\n",
+    "    print(f\"   {name} head: fraction of sampled tokens with |pi/mu - 1| > 0.2 at theta=theta_old: {out:.4f}   Pearson corr(p_train, p_sampler) = {corr:.6f}\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "   bf16 head: fraction of sampled tokens with |pi/mu - 1| > 0.2 at theta=theta_old: 0.0000   Pearson corr(p_train, p_sampler) = 0.999793\n",
+      "   fp32 head: fraction of sampled tokens with |pi/mu - 1| > 0.2 at theta=theta_old: 0.0000   Pearson corr(p_train, p_sampler) = 1.000000\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7. Naive vs stable weight computation under mismatch (log-space)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "source": [
+    "Elw, Vlw, chi2, pi_t, mu_t = per_position_stats(10.0, torch.bfloat16, T=4096)\n",
+    "tokens = torch.stack([torch.multinomial(mu_t[t], 1, replacement=True, generator=g) for t in range(4096)], 0).squeeze(1)\n",
+    "lr_ = (pi_t.log().gather(1, tokens[:, None]) - mu_t.log().gather(1, tokens[:, None])).squeeze(1)\n",
+    "logw = float(lr_.sum())\n",
+    "naive32 = torch.tensor(1.0, dtype=torch.float32)\n",
+    "for v in lr_.exp().to(torch.float32): naive32 = naive32 * v\n",
+    "print(f\"   T=4096 bf16-mismatch sequence: log w = {logw:.2f};  naive fp32 running product = {float(naive32)};  exp(sum log) in fp64 = {math.exp(logw):.3e}\")\n",
+    "check(\"naive fp32 product vs log-space: agree only when representable (log w within fp32 range)\", 0.0 if (abs(logw) > 80 or abs(float(naive32) - math.exp(logw)) < 1e-4 * math.exp(logw)) else 1.0, 0.0)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "   T=4096 bf16-mismatch sequence: log w = -1.55;  naive fp32 running product = 0.21136480569839478;  exp(sum log) in fp64 = 2.114e-01\n",
+      "CHECK|exact|naive fp32 product vs log-space: agree only when representable (log w within fp32 range)|0.000e+00|tol=0.0e+00|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 8. Naive vs efficient TIS-corrected gradient estimate (sampled from the bf16 sampler mu)\n",
+    "naive: one backward per sample;  efficient: one backward through the surrogate with DETACHED truncated weights."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "source": [
+    "lm = TinyLM(V=4, L=4, seed=0, scale=4.0)\n",
+    "theta = lm.theta0.clone(); R = default_reward(lm)\n",
+    "pi_ = seq_logprob_rounded(theta, None).exp(); mu_ = seq_logprob_rounded(theta, torch.bfloat16).exp()\n",
+    "bb = float(pi_ @ R)\n",
+    "gen = torch.Generator().manual_seed(6)\n",
+    "idx = torch.multinomial(mu_, 500, replacement=True, generator=gen)\n",
+    "Cc = 1.02\n",
+    "w_ = (pi_ / mu_)[idx]\n",
+    "acc = torch.zeros_like(theta)\n",
+    "for j, i in enumerate(idx.tolist()):\n",
+    "    t = theta.clone().requires_grad_(True)\n",
+    "    acc += float(min(w_[j], Cc)) * float(R[i] - bb) * torch.autograd.grad(seq_logprob_rounded(t, None)[i], t)[0]\n",
+    "g_naive = acc / len(idx)\n",
+    "t = theta.clone().requires_grad_(True)\n",
+    "sur = (torch.clamp(w_, max=Cc).detach() * (R[idx] - bb) * seq_logprob_rounded(t, None)[idx]).mean()\n",
+    "g_eff = torch.autograd.grad(sur, t)[0]\n",
+    "check(\"TIS estimate: per-sample loop == single backward through detached-weight surrogate\", maxabs(g_naive, g_eff), 1e-13)\n",
+    "print(\"done; checks:\", len(CHECKS))"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|TIS estimate: per-sample loop == single backward through detached-weight surrogate|3.123e-16|tol=1.0e-13|PASS\n",
+      "done; checks: 39\n"
+     ]
+    }
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

code/notebooks/0109-kl-estimators.ipynb

@@ -0,0 +1,436 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 0109 - KL regularisation estimators k1, k2, k3: value bias/variance, what their GRADIENTS optimise, why beta = 0\n",
+    "Companion to `notes/0109 - KL Estimators.pdf`.  Samples x ~ q = pi_theta,  p = pi_ref,  r = p/q,  KL(q||p) = E_q[-log r]."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "source": [
+    "import sys, math\n",
+    "sys.path.insert(0, \"..\")\n",
+    "from _m1_lib import *\n",
+    "import numpy as np"
+   ],
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Values.  E_q[k1] = E_q[k3] = KL(q||p);  E_q[k2] = KL + bias,  bias = -sum_{n>=3} E[l^n]/n!,  l = log r\n",
+    "  k1 = -log r,   k2 = (log r)^2 / 2,   k3 = (r-1) - log r        (Schulman, \"Approximating KL divergence\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "source": [
+    "def estimators(logr):\n",
+    "    return -logr, 0.5 * logr ** 2, torch.expm1(logr) - logr\n",
+    "gen = torch.Generator().manual_seed(2)\n",
+    "for trial in range(3):\n",
+    "    K = 6\n",
+    "    q = torch.softmax(torch.randn(K, generator=gen) * 1.2, 0); p = torch.softmax(torch.randn(K, generator=gen) * 1.2, 0)\n",
+    "    logr = (p / q).log()\n",
+    "    kl = float((q * (q / p).log()).sum())\n",
+    "    k1, k2, k3 = estimators(logr)\n",
+    "    check(f\"categorical #{trial}: E_q[k1] = KL(q||p)\", abs(float((q * k1).sum()) - kl), 1e-14)\n",
+    "    check(f\"categorical #{trial}: E_q[k3] = KL(q||p)\", abs(float((q * k3).sum()) - kl), 1e-14)\n",
+    "    bias2 = float((q * k2).sum()) - kl\n",
+    "    series = -sum(float((q * logr ** n).sum()) / math.factorial(n) for n in range(3, 40))\n",
+    "    check(f\"categorical #{trial}: bias(k2) = -sum_(n>=3) E[l^n]/n!  (series, 40 terms)\", abs(bias2 - series), 1e-12)\n",
+    "    # the master identity KL = sum_{n>=2} E[l^n]/n!  (from E_q[e^l] = 1)\n",
+    "    ms = sum(float((q * logr ** n).sum()) / math.factorial(n) for n in range(2, 40))\n",
+    "    check(f\"categorical #{trial}: KL = sum_(n>=2) E_q[l^n]/n!\", abs(ms - kl), 1e-12)\n",
+    "    print(f\"      KL = {kl:.5f}   E[k2] - KL = {bias2:+.5f}   (leading term -E[l^3]/6 = {-float((q*logr**3).sum())/6:+.5f})\")\n",
+    "# pointwise positivity\n",
+    "rg = torch.exp(torch.linspace(-12, 12, 100001))\n",
+    "k1g, k2g, k3g = estimators(rg.log())\n",
+    "check(\"k3 >= 0 and k2 >= 0 pointwise (min over a log-grid of r)\", max(0.0, -float(min(k3g.min(), k2g.min()))), 0.0)\n",
+    "print(f\"   k1 < 0 on r > 1 (fraction of log-grid where k1<0: {float((k1g<0).double().mean()):.3f}) -> k1 is a noisy, sign-changing estimator\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|categorical #0: E_q[k1] = KL(q||p)|2.776e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|categorical #0: E_q[k3] = KL(q||p)|2.914e-16|tol=1.0e-14|PASS\n",
+      "CHECK|exact|categorical #0: bias(k2) = -sum_(n>=3) E[l^n]/n!  (series, 40 terms)|2.810e-16|tol=1.0e-12|PASS\n",
+      "CHECK|exact|categorical #0: KL = sum_(n>=2) E_q[l^n]/n!|2.776e-16|tol=1.0e-12|PASS\n",
+      "      KL = 0.10605   E[k2] - KL = -0.01866   (leading term -E[l^3]/6 = -0.01211)\n",
+      "CHECK|exact|categorical #1: E_q[k1] = KL(q||p)|1.110e-16|tol=1.0e-14|PASS\n",
+      "CHECK|exact|categorical #1: E_q[k3] = KL(q||p)|2.220e-16|tol=1.0e-14|PASS\n",
+      "CHECK|exact|categorical #1: bias(k2) = -sum_(n>=3) E[l^n]/n!  (series, 40 terms)|1.943e-16|tol=1.0e-12|PASS\n",
+      "CHECK|exact|categorical #1: KL = sum_(n>=2) E_q[l^n]/n!|2.220e-16|tol=1.0e-12|PASS\n",
+      "      KL = 0.28620   E[k2] - KL = +0.01383   (leading term -E[l^3]/6 = +0.06153)\n",
+      "CHECK|exact|categorical #2: E_q[k1] = KL(q||p)|0.000e+00|tol=1.0e-14|PASS\n",
+      "CHECK|exact|categorical #2: E_q[k3] = KL(q||p)|0.000e+00|tol=1.0e-14|PASS\n",
+      "CHECK|exact|categorical #2: bias(k2) = -sum_(n>=3) E[l^n]/n!  (series, 40 terms)|5.551e-17|tol=1.0e-12|PASS\n",
+      "CHECK|exact|categorical #2: KL = sum_(n>=2) E_q[l^n]/n!|1.110e-16|tol=1.0e-12|PASS\n",
+      "      KL = 0.61184   E[k2] - KL = +0.15474   (leading term -E[l^3]/6 = +0.31903)\n",
+      "CHECK|exact|k3 >= 0 and k2 >= 0 pointwise (min over a log-grid of r)|0.000e+00|tol=0.0e+00|PASS\n",
+      "   k1 < 0 on r > 1 (fraction of log-grid where k1<0: 0.500) -> k1 is a noisy, sign-changing estimator\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Gaussian closed forms (equal variance, shift mu => KL = mu^2/2).  With s2 = mu^2:\n",
+    "  Var(k1) = s2,   Var(k2) = s2^3/4 + s2^2/2,   Var(k3) = e^{s2} - 1 - s2,   bias(k2) = s2^2/8.   Verified by Gauss-Hermite quadrature."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "source": [
+    "xs, ws = np.polynomial.hermite_e.hermegauss(120); ws = ws / ws.sum()\n",
+    "def gauss_moments(shift):\n",
+    "    x = torch.tensor(xs); w = torch.tensor(ws)\n",
+    "    logr = shift * x - 0.5 * shift ** 2             # q = N(0,1), p = N(shift,1): log p/q = shift*x - shift^2/2  (E_q[r] = 1)\n",
+    "    k = estimators(logr)\n",
+    "    kl = shift ** 2 / 2\n",
+    "    out = []\n",
+    "    for kk in k:\n",
+    "        m = float((w * kk).sum()); v = float((w * kk ** 2).sum()) - m ** 2\n",
+    "        out.append((m, v))\n",
+    "    return kl, out\n",
+    "print(\"   KL     est   mean-KL (bias)     std/KL      closed-form std/KL\")\n",
+    "for shift in (0.1, math.sqrt(2 * 0.05), 1.0):\n",
+    "    kl, out = gauss_moments(shift)\n",
+    "    s2 = shift ** 2\n",
+    "    cf = [s2, s2 ** 3 / 4 + s2 ** 2 / 2, math.exp(s2) - 1 - s2]\n",
+    "    cb = [0.0, s2 ** 2 / 8, 0.0]\n",
+    "    for i, nm in enumerate((\"k1\", \"k2\", \"k3\")):\n",
+    "        m, v = out[i]\n",
+    "        print(f\"   {kl:.4f}  {nm}  {m-kl:+.3e}   {math.sqrt(v)/kl:9.4f}   {math.sqrt(cf[i])/kl:9.4f}\")\n",
+    "        check(f\"Gaussian shift={shift:.3f}: {nm} bias == closed form\", abs((m - kl) - cb[i]), 1e-12)\n",
+    "        check(f\"Gaussian shift={shift:.3f}: {nm} variance == closed form\", abs(v - cf[i]) / cf[i], 1e-10)\n",
+    "# compare with the table in Schulman's note (KL=0.005 and KL=0.5): std/KL = 20, 1.42, 1.42 and 2, 1.73, 1.7 ; bias(k2)/KL = 0.002, 0.25\n",
+    "kl, out = gauss_moments(0.1)\n",
+    "check(\"matches Schulman's table at KL=0.005: std(k1)/KL ~ 20\", abs(math.sqrt(out[0][1]) / kl - 20.0), 0.2)\n",
+    "kl, out = gauss_moments(1.0)\n",
+    "check(\"matches Schulman's table at KL=0.5: bias(k2)/KL = 0.25\", abs((out[1][0] - kl) / kl - 0.25), 1e-9)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "   KL     est   mean-KL (bias)     std/KL      closed-form std/KL\n",
+      "   0.0050  k1  +1.735e-18     20.0000     20.0000\n",
+      "CHECK|exact|Gaussian shift=0.100: k1 bias == closed form|1.735e-18|tol=1.0e-12|PASS\n",
+      "CHECK|exact|Gaussian shift=0.100: k1 variance == closed form|1.735e-16|tol=1.0e-10|PASS\n",
+      "   0.0050  k2  +1.250e-05      1.4177      1.4177\n",
+      "CHECK|exact|Gaussian shift=0.100: k2 bias == closed form|1.220e-18|tol=1.0e-12|PASS\n",
+      "CHECK|exact|Gaussian shift=0.100: k2 variance == closed form|1.349e-16|tol=1.0e-10|PASS\n",
+      "   0.0050  k3  +2.602e-18      1.4166      1.4166\n",
+      "CHECK|exact|Gaussian shift=0.100: k3 bias == closed form|2.602e-18|tol=1.0e-12|PASS\n",
+      "CHECK|exact|Gaussian shift=0.100: k3 variance == closed form|2.200e-12|tol=1.0e-10|PASS\n",
+      "   0.0500  k1  +1.388e-17      6.3246      6.3246\n",
+      "CHECK|exact|Gaussian shift=0.316: k1 bias == closed form|1.388e-17|tol=1.0e-12|PASS\n",
+      "CHECK|exact|Gaussian shift=0.316: k1 variance == closed form|2.776e-16|tol=1.0e-10|PASS\n",
+      "   0.0500  k2  +1.250e-03      1.4491      1.4491\n",
+      "CHECK|exact|Gaussian shift=0.316: k2 bias == closed form|1.475e-17|tol=1.0e-12|PASS\n",
+      "CHECK|exact|Gaussian shift=0.316: k2 variance == closed form|0.000e+00|tol=1.0e-10|PASS\n",
+      "   0.0500  k3  +1.388e-17      1.4382      1.4382\n",
+      "CHECK|exact|Gaussian shift=0.316: k3 bias == closed form|1.388e-17|tol=1.0e-12|PASS\n",
+      "CHECK|exact|Gaussian shift=0.316: k3 variance == closed form|1.543e-14|tol=1.0e-10|PASS\n",
+      "   0.5000  k1  +0.000e+00      2.0000      2.0000\n",
+      "CHECK|exact|Gaussian shift=1.000: k1 bias == closed form|0.000e+00|tol=1.0e-12|PASS\n",
+      "CHECK|exact|Gaussian shift=1.000: k1 variance == closed form|2.220e-16|tol=1.0e-10|PASS\n",
+      "   0.5000  k2  +1.250e-01      1.7321      1.7321\n",
+      "CHECK|exact|Gaussian shift=1.000: k2 bias == closed form|1.110e-16|tol=1.0e-12|PASS\n",
+      "CHECK|exact|Gaussian shift=1.000: k2 variance == closed form|5.921e-16|tol=1.0e-10|PASS\n",
+      "   0.5000  k3  +1.110e-16      1.6950      1.6950\n",
+      "CHECK|exact|Gaussian shift=1.000: k3 bias == closed form|1.110e-16|tol=1.0e-12|PASS\n",
+      "CHECK|exact|Gaussian shift=1.000: k3 variance == closed form|1.546e-16|tol=1.0e-10|PASS\n",
+      "CHECK|exact|matches Schulman's table at KL=0.005: std(k1)/KL ~ 20|0.000e+00|tol=2.0e-01|PASS\n",
+      "CHECK|exact|matches Schulman's table at KL=0.5: bias(k2)/KL = 0.25|2.220e-16|tol=1.0e-09|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Heavy tails: Var(k3) is infinite when chi^2(p||q) is; one token the reference loves and the policy nearly never samples"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "source": [
+    "q = torch.tensor([0.001, 0.999]); p = torch.tensor([0.30, 0.70])\n",
+    "logr = (p / q).log(); k1, k2, k3 = estimators(logr)\n",
+    "kl = float((q * (q / p).log()).sum())\n",
+    "for nm, k in ((\"k1\", k1), (\"k2\", k2), (\"k3\", k3)):\n",
+    "    m = float((q * k).sum()); v = float((q * k ** 2).sum()) - m ** 2\n",
+    "    print(f\"   {nm}: mean {m:+.5f}  std {math.sqrt(v):.5f}   (KL = {kl:.5f})   max value {float(k.max()):.1f}\")\n",
+    "chi2_pq = float((p ** 2 / q).sum() - 1)\n",
+    "print(f\"   chi2(p||q) = E_q[r^2]-1 = {chi2_pq:.1f}  -> Var(k3) is dominated by one rare token with r = {float((p/q)[0]):.0f}\")\n",
+    "check(\"E_q[r^2] - 1 = chi2(p||q) exactly\", abs(float((q * (p / q) ** 2).sum()) - 1 - chi2_pq), 1e-9)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "   k1: mean +0.34961  std 0.19152   (KL = 0.34961)   max value 0.4\n",
+      "   k2: mean +0.07946  std 0.51214   (KL = 0.34961)   max value 16.3\n",
+      "   k3: mean +0.34961  std 9.26842   (KL = 0.34961)   max value 293.3\n",
+      "   chi2(p||q) = E_q[r^2]-1 = 89.5  -> Var(k3) is dominated by one rare token with r = 300\n",
+      "CHECK|exact|E_q[r^2] - 1 = chi2(p||q) exactly|0.000e+00|tol=1.0e-09|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. GRADIENTS (on-policy, single-step categorical policy with logits theta; samples x ~ q_theta DETACHED, gradient flows through q_theta in r)\n",
+    "  k1-as-loss:  E[grad k1] = E[grad log q] = 0                        -> useless as a regulariser\n",
+    "  k2-as-loss:  E[grad k2] = E[(log q - log p) grad log q] = grad KL(q||p)           (the reverse KL we meant)\n",
+    "  k3-as-loss:  E[grad k3] = E[(1 - r) grad log q]        = -sum_x p grad log q = grad KL(p||q)   (the FORWARD KL)\n",
+    "  k1-in-reward (detached coefficient):  E[(log q - log p) grad log q] = grad KL(q||p)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "source": [
+    "gen = torch.Generator().manual_seed(7)\n",
+    "K = 6\n",
+    "th_ref = torch.randn(K, generator=gen) * 1.0\n",
+    "def grads(delta):\n",
+    "    th = (th_ref + delta * torch.randn(K, generator=gen)).clone().requires_grad_(True)\n",
+    "    q = torch.softmax(th, 0); p = torch.softmax(th_ref, 0)\n",
+    "    lq, lp = torch.log_softmax(th, 0), torch.log_softmax(th_ref, 0)\n",
+    "    qd = q.detach()\n",
+    "    r = (lp - lq).exp()\n",
+    "    # expectation over x ~ q (weights detached) of the gradient of each estimator\n",
+    "    g = {}\n",
+    "    g['k1'] = torch.autograd.grad((qd * (-(lp - lq))).sum(), th, retain_graph=True)[0]\n",
+    "    g['k2'] = torch.autograd.grad((qd * 0.5 * (lp - lq) ** 2).sum(), th, retain_graph=True)[0]\n",
+    "    g['k3'] = torch.autograd.grad((qd * (r - 1 - (lp - lq))).sum(), th, retain_graph=True)[0]\n",
+    "    # k1 in reward: detached coefficient (log q - log p) times score\n",
+    "    g['k1r'] = torch.autograd.grad((qd * (lq - lp).detach() * lq).sum(), th, retain_graph=True)[0]\n",
+    "    g['rev'] = torch.autograd.grad((q * (lq - lp)).sum(), th, retain_graph=True)[0]          # grad KL(q||p), exact\n",
+    "    g['fwd'] = torch.autograd.grad((p * (lp - lq)).sum(), th, retain_graph=True)[0]          # grad KL(p||q), exact\n",
+    "    return g\n",
+    "g = grads(0.8)\n",
+    "check(\"E[grad k1] = 0  (k1 as a loss carries no regularising gradient)\", float(g['k1'].abs().max()), 1e-14)\n",
+    "check(\"E[grad k2] == grad KL(q||p)  (reverse KL)\", maxabs(g['k2'], g['rev']), 1e-14)\n",
+    "check(\"k1-in-reward (detached coefficient) == grad KL(q||p)\", maxabs(g['k1r'], g['rev']), 1e-14)\n",
+    "check(\"E[grad k3] == grad KL(p||q)  (FORWARD KL, not the one k3 estimates)\", maxabs(g['k3'], g['fwd']), 1e-14)\n",
+    "print(f\"   ||grad KL(q||p) - grad KL(p||q)|| / ||grad KL(q||p)|| = {float((g['rev']-g['fwd']).norm()/g['rev'].norm()):.3f}   (different vectors)\")\n",
+    "print(\"   relative gap between the two gradients as q -> p (first-order agreement, 'k3 is a first-order approximation'):\")\n",
+    "prev = None\n",
+    "for delta in (0.4, 0.2, 0.1, 0.05, 0.025):\n",
+    "    gen.manual_seed(7); _ = torch.randn(K, generator=gen)          # same direction each time\n",
+    "    gg = grads(delta)\n",
+    "    gap = float((gg['rev'] - gg['fwd']).norm() / gg['rev'].norm())\n",
+    "    print(f\"     delta={delta:5.3f}   gap = {gap:.4f}\")\n",
+    "    prev = gap\n",
+    "assert prev < 0.05\n",
+    "check(\"gradient gap forward vs reverse KL vanishes as delta -> 0 (gap at delta=0.025 < 5%)\", 0.0, 0.0)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|E[grad k1] = 0  (k1 as a loss carries no regularising gradient)|5.551e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|E[grad k2] == grad KL(q||p)  (reverse KL)|5.551e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|k1-in-reward (detached coefficient) == grad KL(q||p)|5.551e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|E[grad k3] == grad KL(p||q)  (FORWARD KL, not the one k3 estimates)|8.327e-17|tol=1.0e-14|PASS\n",
+      "   ||grad KL(q||p) - grad KL(p||q)|| / ||grad KL(q||p)|| = 0.215   (different vectors)\n",
+      "   relative gap between the two gradients as q -> p (first-order agreement, 'k3 is a first-order approximation'):\n",
+      "     delta=0.400   gap = 0.1533\n",
+      "     delta=0.200   gap = 0.0922\n",
+      "     delta=0.100   gap = 0.0505\n",
+      "     delta=0.050   gap = 0.0264\n",
+      "     delta=0.025   gap = 0.0135\n",
+      "CHECK|exact|gradient gap forward vs reverse KL vanishes as delta -> 0 (gap at delta=0.025 < 5%)|0.000e+00|tol=0.0e+00|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. SEQUENCES: per-token k2/k3 'as loss' ignore that earlier tokens change the later contexts\n",
+    "  true:   grad KL_seq(pi||ref) = E[ (sum_t l_t) * score(o) ] = E[ sum_t (sum_{t'>=t} l_t') s_t ]   (\"k1 in reward\" with reward-to-go)\n",
+    "  per-token k2 loss, summed over t:   E[ sum_t l_t s_t ]  =  true  -  sum_t E[ KL_t(s_t) * score(o_<t) ]    (missing prefix-score term)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "source": [
+    "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n",
+    "th = lm.theta0.clone(); th_r = perturb(th, 1.0, seed=2)\n",
+    "pi = lm.probs(th)\n",
+    "lpi, lref = lm.token_logprobs(th), lm.token_logprobs(th_r)\n",
+    "l_tok = (lpi - lref) * lm.mask                                     # l_t = log pi - log ref  = -log r_t\n",
+    "Ttens = token_score_tensor(lm, th)                                 # [N, L, P]\n",
+    "S = Ttens.sum(1)\n",
+    "KL_seq = float((pi * l_tok.sum(1)).sum())\n",
+    "t_ = th.clone().requires_grad_(True)\n",
+    "KLt = (lm.probs(t_) * (lm.seq_logprob(t_) - lm.seq_logprob(th_r))).sum()\n",
+    "g_true = torch.autograd.grad(KLt, t_)[0]\n",
+    "g_seq = (pi[:, None] * l_tok.sum(1)[:, None] * S).sum(0)\n",
+    "rtg = l_tok.flip(1).cumsum(1).flip(1) * lm.mask\n",
+    "g_rtg = (pi[:, None] * (Ttens * rtg[:, :, None]).sum(1)).sum(0)\n",
+    "g_tok = (pi[:, None] * (Ttens * l_tok[:, :, None]).sum(1)).sum(0)                  # per-token k2-as-loss gradient, summed over t\n",
+    "check(\"E[(sum_t l_t) * score] == grad KL_seq(pi||ref)\", maxabs(g_seq, g_true), 1e-14)\n",
+    "check(\"reward-to-go form E[sum_t (sum_{t'>=t} l_t') s_t] == grad KL_seq\", maxabs(g_rtg, g_true), 1e-14)\n",
+    "print(f\"   per-token k2-as-loss gradient vs true sequence-KL gradient: relative error {float((g_tok-g_true).norm()/g_true.norm()):.4f}\")\n",
+    "# the missing term: sum_t E[ KL_t(s_t) * score(o_<t) ] where KL_t(s) = KL(pi(.|s) || ref(.|s)) and score(o_<t) = sum_{s<t} s_s\n",
+    "prefix_score = torch.cumsum(Ttens, 1) - Ttens                                      # sum_{s<t} s_s\n",
+    "KLt_tok = torch.zeros(lm.N, lm.L)\n",
+    "for i, sq in enumerate(lm.seqs):\n",
+    "    for t in range(len(sq)):\n",
+    "        cp, cr = lm.cond_probs(th, sq[:t]), lm.cond_probs(th_r, sq[:t])\n",
+    "        KLt_tok[i, t] = (cp * (cp.log() - cr.log())).sum()\n",
+    "missing = (pi[:, None] * (KLt_tok[:, :, None] * prefix_score).sum(1)).sum(0)\n",
+    "check(\"per-token k2 loss gradient == true gradient - sum_t E[KL_t * prefix score]  (exact identity)\", maxabs(g_tok, g_true - missing), 1e-13)\n",
+    "# per-token k3 loss gradient equals the gradient of the sum of FORWARD per-state KLs with frozen context distribution\n",
+    "r_tok = (lref - lpi).exp() * lm.mask + (1 - lm.mask)\n",
+    "g_k3 = (pi[:, None] * (Ttens * ((1 - r_tok) * lm.mask)[:, :, None]).sum(1)).sum(0)\n",
+    "t_ = th.clone().requires_grad_(True)\n",
+    "tot = 0.0\n",
+    "for s in lm.prefixes():\n",
+    "    pr = 1.0\n",
+    "    for k in range(len(s)): pr *= float(lm.cond_probs(th, s[:k])[s[k]])           # P_pi(context = s) (detached)\n",
+    "    cp_t, cr = lm.cond_probs(t_, s), lm.cond_probs(th_r, s)\n",
+    "    tot = tot + pr * (cr * (cr.log() - cp_t.log())).sum()                        # KL(ref(.|s) || pi_theta(.|s)) with frozen context weight\n",
+    "g_fwd_frozen = torch.autograd.grad(tot, t_)[0]\n",
+    "check(\"per-token k3-as-loss gradient == grad of sum_s P_pi(s) KL(ref(.|s)||pi_theta(.|s)) (context weights frozen)\", maxabs(g_k3, g_fwd_frozen), 1e-13)\n",
+    "# values: sequence-level KL from the three per-token estimators\n",
+    "k1_seq = float((pi * l_tok.sum(1)).sum())\n",
+    "k3_tok = (torch.expm1(lref - lpi) - (lref - lpi)) * lm.mask\n",
+    "k2_tok = 0.5 * (lpi - lref) ** 2 * lm.mask\n",
+    "print(f\"   sequence KL = {KL_seq:.5f};  E[sum_t k1] = {k1_seq:.5f};  E[sum_t k3] = {float((pi*k3_tok.sum(1)).sum()):.5f};  E[sum_t k2] = {float((pi*k2_tok.sum(1)).sum()):.5f}\")\n",
+    "check(\"E[sum_t k3_t] = KL_seq (per-token estimator is unbiased for the sequence KL value)\", abs(float((pi * k3_tok.sum(1)).sum()) - KL_seq), 1e-13)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|E[(sum_t l_t) * score] == grad KL_seq(pi||ref)|7.633e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|reward-to-go form E[sum_t (sum_{t'>=t} l_t') s_t] == grad KL_seq|8.327e-17|tol=1.0e-14|PASS\n",
+      "   per-token k2-as-loss gradient vs true sequence-KL gradient: relative error 0.0376\n",
+      "CHECK|exact|per-token k2 loss gradient == true gradient - sum_t E[KL_t * prefix score]  (exact identity)|9.714e-17|tol=1.0e-13|PASS\n",
+      "CHECK|exact|per-token k3-as-loss gradient == grad of sum_s P_pi(s) KL(ref(.|s)||pi_theta(.|s)) (context weights frozen)|5.898e-17|tol=1.0e-13|PASS\n",
+      "   sequence KL = 0.01885;  E[sum_t k1] = 0.01885;  E[sum_t k3] = 0.01885;  E[sum_t k2] = 0.02103\n",
+      "CHECK|exact|E[sum_t k3_t] = KL_seq (per-token estimator is unbiased for the sequence KL value)|6.245e-17|tol=1.0e-13|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Naive vs stable evaluation of k3 (fp32): (r-1)-log r cancels catastrophically when |log r| is small"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "source": [
+    "lr32 = torch.tensor([1e-1, 1e-2, 1e-3, 1e-4, 1e-5], dtype=torch.float32)\n",
+    "naive = (lr32.exp() - 1) - lr32\n",
+    "stable = torch.expm1(lr32) - lr32\n",
+    "truth = (torch.expm1(lr32.double()) - lr32.double())\n",
+    "print(\"   log r     naive fp32        expm1 fp32       fp64 truth       naive rel.err   stable rel.err\")\n",
+    "for a, n_, s_, t_ in zip(lr32.tolist(), naive.tolist(), stable.tolist(), truth.tolist()):\n",
+    "    print(f\"   {a:7.0e}  {n_:+.4e}   {s_:+.4e}   {t_:+.4e}    {abs(n_-t_)/t_:9.2e}     {abs(s_-t_)/t_:9.2e}\")\n",
+    "check(\"expm1 form: relative error < 1e-3 for |log r| >= 1e-4 (fp32)\", float(((stable.double() - truth).abs() / truth)[:4].max()), 1e-3)\n",
+    "assert float(((naive.double() - truth).abs() / truth)[3]) > 0.1\n",
+    "check(\"naive form is wrong by >10% at |log r| = 1e-4 (fp32)\", 0.0, 0.0)\n",
+    "# The expm1 form removes the cancellation inside exp(x)-1 but NOT the one in expm1(x) - x ~ x^2/2: its worst-case relative error is ~ 2u/|x|.\n",
+    "# Worst case over a dense random grid per decade (fp32, both signs), vs the 2u/|x|_min law; and a series fallback for small |x|.\n",
+    "u32 = 2.0 ** -24\n",
+    "gen_n = torch.Generator().manual_seed(11)\n",
+    "def k3_expm1(x): return torch.expm1(x) - x\n",
+    "def k3_hybrid(x):\n",
+    "    ser = x * x / 2 * (1 + x / 3 + x * x / 12)                              # Taylor of expm1(x)-x, error ~ x^5/120 relative x^3/60\n",
+    "    return torch.where(x.abs() < 1e-2, ser, k3_expm1(x))\n",
+    "print(\"   |x| decade        worst rel.err expm1 form   2u/|x|min     worst rel.err hybrid (series for |x|<1e-2)\")\n",
+    "worst_ratio, worst_hybrid = 0.0, 0.0\n",
+    "for lo in (1e-1, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6):\n",
+    "    hi = 10 * lo\n",
+    "    xg = torch.exp(torch.rand(400_000, generator=gen_n, dtype=torch.float64) * math.log(10) + math.log(lo)).float()\n",
+    "    xg = torch.cat([xg, -xg]); tr = torch.expm1(xg.double()) - xg.double()\n",
+    "    re = float(((k3_expm1(xg).double() - tr).abs() / tr).max()); rh = float(((k3_hybrid(xg).double() - tr).abs() / tr).max())\n",
+    "    law = 2 * u32 / lo\n",
+    "    worst_ratio = max(worst_ratio, re / law); worst_hybrid = max(worst_hybrid, rh)\n",
+    "    print(f\"   [{lo:.0e},{hi:.0e}]      {re:12.3e}              {law:10.2e}     {rh:12.3e}\")\n",
+    "check(\"expm1 form: worst-case relative error <= 2u/|x| (max ratio to the law)\", worst_ratio, 1.0)\n",
+    "check(\"hybrid (Taylor for |x|<1e-2, expm1 form above): worst relative error over |x| in [1e-6, 1)\", worst_hybrid, 1e-4)\n",
+    "print(\"done; checks:\", len(CHECKS))"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "   log r     naive fp32        expm1 fp32       fp64 truth       naive rel.err   stable rel.err\n",
+      "     1e-01  +5.1710e-03   +5.1709e-03   +5.1709e-03     8.79e-06      1.49e-07\n",
+      "     1e-02  +5.0178e-05   +5.0167e-05   +5.0167e-05     2.14e-04      9.17e-06\n",
+      "     1e-03  +5.2352e-07   +5.0012e-07   +5.0017e-07     4.67e-02      9.30e-05\n",
+      "     1e-04  +1.6596e-08   +4.9986e-09   +5.0002e-09     2.32e+00      3.17e-04\n",
+      "     1e-05  +1.3581e-08   +5.0022e-11   +5.0000e-11     2.71e+02      4.41e-04\n",
+      "CHECK|exact|expm1 form: relative error < 1e-3 for |log r| >= 1e-4 (fp32)|3.167e-04|tol=1.0e-03|PASS\n",
+      "CHECK|exact|naive form is wrong by >10% at |log r| = 1e-4 (fp32)|0.000e+00|tol=0.0e+00|PASS\n",
+      "   |x| decade        worst rel.err expm1 form   2u/|x|min     worst rel.err hybrid (series for |x|<1e-2)\n",
+      "   [1e-01,1e+00]         1.033e-06                1.19e-06        1.033e-06\n",
+      "   [1e-02,1e-01]         9.291e-06                1.19e-05        9.291e-06\n",
+      "   [1e-03,1e-02]         1.161e-04                1.19e-04        2.203e-07\n",
+      "   [1e-04,1e-03]         9.735e-04                1.19e-03        2.035e-07\n",
+      "   [1e-05,1e-04]         9.034e-03                1.19e-02        1.667e-07\n",
+      "   [1e-06,1e-05]         1.150e-01                1.19e-01        1.646e-07\n",
+      "CHECK|exact|expm1 form: worst-case relative error <= 2u/|x| (max ratio to the law)|9.736e-01|tol=1.0e+00|PASS\n",
+      "CHECK|exact|hybrid (Taylor for |x|<1e-2, expm1 form above): worst relative error over |x| in [1e-6, 1)|9.291e-06|tol=1.0e-04|PASS\n",
+      "done; checks: 48\n"
+     ]
+    }
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

code/notebooks/0110-reward-shaping-verifier.ipynb

@@ -0,0 +1,417 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 0110 - Reward shaping for verifier rewards: binary rewards, dead groups, overlong penalties as reward, verifier errors\n",
+    "Companion to `notes/0110 - Reward Shaping for Verifier Rewards.pdf`.  fp64; exact enumeration on the 121-sequence toy LM."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "source": [
+    "import sys, math, itertools\n",
+    "sys.path.insert(0, \"..\")\n",
+    "from _m1_lib import *\n",
+    "import numpy as np\n",
+    "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n",
+    "theta = lm.theta0.clone()\n",
+    "R = default_reward(lm)\n",
+    "p_seq = lm.probs(theta)\n",
+    "S = score_matrix(lm, theta)\n",
+    "N = lm.N\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "pr = (lm.probs(th) * R).sum()\n",
+    "grad_p = torch.autograd.grad(pr, th)[0]\n",
+    "p = float(pr.detach())"
+   ],
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Binary reward: grad p = p (1-p) ( E[score | success] - E[score | failure] )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "source": [
+    "m1 = (p_seq * R)[:, None] * S\n",
+    "m0 = (p_seq * (1 - R))[:, None] * S\n",
+    "E_s_succ = m1.sum(0) / p; E_s_fail = m0.sum(0) / (1 - p)\n",
+    "check(\"grad p = p (1-p) (E[s|succ] - E[s|fail])\", maxabs(grad_p, p * (1 - p) * (E_s_succ - E_s_fail)), 1e-14)\n",
+    "check(\"grad p = p E[s | success]  (rejection-sampling / expert-iteration form)\", maxabs(grad_p, p * E_s_succ), 1e-14)\n",
+    "print(f\"   p = {p:.4f};  signal scale p(1-p) = {p*(1-p):.4f}\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|grad p = p (1-p) (E[s|succ] - E[s|fail])|6.939e-18|tol=1.0e-14|PASS\n",
+      "CHECK|exact|grad p = p E[s | success]  (rejection-sampling / expert-iteration form)|1.388e-17|tol=1.0e-14|PASS\n",
+      "   p = 0.1659;  signal scale p(1-p) = 0.1384\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Dead groups.  Binary bandit with logit parameter: the group estimators collapse to functions of p_hat = k/G\n",
+    "  centred (Dr.GRPO):  g = p_hat (1 - p_hat)            (plug-in variance of the rewards);   std-normalised (GRPO):  g = sqrt(p_hat (1 - p_hat))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "source": [
+    "def bern_group(G, pp):\n",
+    "    ks = np.arange(G + 1)\n",
+    "    w = np.array([math.comb(G, k) * pp ** k * (1 - pp) ** (G - k) for k in ks])\n",
+    "    ph = ks / G\n",
+    "    return ks, w, ph\n",
+    "def direct_estimator(G, k, pp, kind):\n",
+    "    \"\"\"(1/G) sum_i A_i * s_i  with s = d log pi / d logit:  success -> (1-p), failure -> -p   (computed from the definitions)\"\"\"\n",
+    "    Rg = torch.cat([torch.ones(k), torch.zeros(G - k)])\n",
+    "    mu = Rg.mean(); sd = Rg.std(unbiased=False)\n",
+    "    A = Rg - mu if kind == 'centre' else (torch.zeros_like(Rg) if float(sd) == 0 else (Rg - mu) / sd)\n",
+    "    s = torch.where(Rg > 0, torch.tensor(1 - pp), torch.tensor(-pp))\n",
+    "    return float((A * s).mean())\n",
+    "worst = 0.0\n",
+    "for G in (2, 4, 16):\n",
+    "    for pp in (0.05, 0.3, 0.7):\n",
+    "        ks, w, ph = bern_group(G, pp)\n",
+    "        for k in ks:\n",
+    "            worst = max(worst, abs(direct_estimator(G, int(k), pp, 'centre') - ph[k] * (1 - ph[k])),\n",
+    "                        abs(direct_estimator(G, int(k), pp, 'std') - math.sqrt(ph[k] * (1 - ph[k]))))\n",
+    "check(\"bandit: centred estimator == p_hat(1-p_hat), std-normalised == sqrt(p_hat(1-p_hat)) (all k, 9 settings)\", worst, 1e-14)\n",
+    "# exact mean / variance / SNR of the estimator as a function of p\n",
+    "print(\"   G=16:   p     P(dead)    E[g_centred]   (1-1/G)p(1-p)   SNR=E^2/Var    E[g_std]/(p(1-p)) vs c_G(p)\")\n",
+    "def c_G_formula(G, pp):          # from note 0105\n",
+    "    bp = lambda n, k: math.comb(n, k) * pp ** k * (1 - pp) ** (n - k)\n",
+    "    a1 = sum(bp(G - 1, K) * (math.sqrt((G - (1 + K)) / (1 + K)) if 1 + K < G else 0.0) for K in range(G))\n",
+    "    a0 = sum(bp(G - 1, K) * (-math.sqrt(K / (G - K)) if 1 <= K < G else 0.0) for K in range(G))\n",
+    "    return a1 - a0\n",
+    "G = 16\n",
+    "snr = {}\n",
+    "for pp in (0.01, 0.05, 0.2, 0.5, 0.8, 0.95, 0.99):\n",
+    "    ks, w, ph = bern_group(G, pp)\n",
+    "    g = ph * (1 - ph); gs = np.sqrt(ph * (1 - ph))\n",
+    "    mean, var = float((w * g).sum()), float((w * g ** 2).sum() - (w * g).sum() ** 2)\n",
+    "    pdead = pp ** G + (1 - pp) ** G\n",
+    "    mean_std = float((w * gs).sum())\n",
+    "    snr[pp] = mean ** 2 / var\n",
+    "    print(f\"         {pp:5.2f}   {pdead:7.4f}   {mean:10.6f}     {(1-1/G)*pp*(1-pp):10.6f}     {snr[pp]:8.3f}     {mean_std/(pp*(1-pp)):7.4f} vs {c_G_formula(G, pp):7.4f}\")\n",
+    "    check(f\"E[g_centred] = (1-1/G) p(1-p)  (p={pp})\", abs(mean - (1 - 1 / G) * pp * (1 - pp)), 1e-15)\n",
+    "    check(f\"E[sqrt(p_hat(1-p_hat))] / (p(1-p)) == c_G(p) of note 0105  (p={pp})\", abs(mean_std / (pp * (1 - pp)) - c_G_formula(G, pp)), 1e-12)\n",
+    "    check(f\"P(dead) = p^G + (1-p)^G = mass of k in {{0,G}}  (p={pp})\", abs(pdead - (w[0] + w[-1])), 1e-15)\n",
+    "assert snr[0.5] > snr[0.2] > snr[0.05] > snr[0.01] and snr[0.5] > snr[0.8] > snr[0.95] > snr[0.99]\n",
+    "check(\"SNR of the group estimator peaks at p = 1/2 and decays towards both ends (monotone on each side)\", 0.0, 0.0)\n",
+    "# dead groups carry exactly zero gradient for every baseline estimator; none-baseline REINFORCE does not\n",
+    "Gs = 4\n",
+    "dead_succ = [int(np.argmax(R.numpy()))] * 0 + [i for i in range(N) if R[i] == 1][:Gs]          # an all-correct group (4 successful sequences)\n",
+    "dead_fail = [i for i in range(N) if R[i] == 0][:Gs]\n",
+    "def group_gradient(idx, kind):\n",
+    "    Rg = R[idx]\n",
+    "    if kind == 'centre': A = Rg - Rg.mean()\n",
+    "    elif kind == 'std':\n",
+    "        sd = Rg.std(unbiased=False); A = torch.zeros_like(Rg) if float(sd) == 0 else (Rg - Rg.mean()) / sd\n",
+    "    elif kind == 'rloo': A = Rg - (Rg.sum() - Rg) / (len(Rg) - 1)\n",
+    "    elif kind == 'plain01': A = Rg                                  # REINFORCE, no baseline, R in {0,1}\n",
+    "    elif kind == 'plainpm': A = 2 * Rg - 1                           # REINFORCE, no baseline, R in {-1,+1}\n",
+    "    return (A[:, None] * S[idx]).mean(0)\n",
+    "print(\"   ||group gradient|| for dead groups (4 samples):\")\n",
+    "for kind in ('centre', 'std', 'rloo', 'plain01', 'plainpm'):\n",
+    "    print(f\"     {kind:8s}   all-correct: {float(group_gradient(dead_succ, kind).norm()):.4f}    all-wrong: {float(group_gradient(dead_fail, kind).norm()):.4f}\")\n",
+    "for kind in ('centre', 'std', 'rloo'):\n",
+    "    check(f\"dead groups have exactly zero gradient ({kind})\", max(float(group_gradient(dead_succ, kind).abs().max()), float(group_gradient(dead_fail, kind).abs().max())), 0.0)\n",
+    "assert float(group_gradient(dead_succ, 'plain01').norm()) > 1e-3 and float(group_gradient(dead_fail, 'plainpm').norm()) > 1e-3\n",
+    "check(\"without a baseline an all-correct (R in {0,1}) or all-wrong (R in {-1,1}) group still moves the policy\", 0.0, 0.0)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|bandit: centred estimator == p_hat(1-p_hat), std-normalised == sqrt(p_hat(1-p_hat)) (all k, 9 settings)|1.665e-16|tol=1.0e-14|PASS\n",
+      "   G=16:   p     P(dead)    E[g_centred]   (1-1/G)p(1-p)   SNR=E^2/Var    E[g_std]/(p(1-p)) vs c_G(p)\n",
+      "          0.01    0.8515     0.009281       0.009281        0.165      3.7330 vs  3.7330\n",
+      "CHECK|exact|E[g_centred] = (1-1/G) p(1-p)  (p=0.01)|0.000e+00|tol=1.0e-15|PASS\n",
+      "CHECK|exact|E[sqrt(p_hat(1-p_hat))] / (p(1-p)) == c_G(p) of note 0105  (p=0.01)|8.882e-16|tol=1.0e-12|PASS\n",
+      "CHECK|exact|P(dead) = p^G + (1-p)^G = mass of k in {0,G}  (p=0.01)|0.000e+00|tol=1.0e-15|PASS\n",
+      "          0.05    0.4401     0.044531       0.044531        0.931      3.2671 vs  3.2671\n",
+      "CHECK|exact|E[g_centred] = (1-1/G) p(1-p)  (p=0.05)|1.388e-17|tol=1.0e-15|PASS\n",
+      "CHECK|exact|E[sqrt(p_hat(1-p_hat))] / (p(1-p)) == c_G(p) of note 0105  (p=0.05)|4.441e-16|tol=1.0e-12|PASS\n",
+      "CHECK|exact|P(dead) = p^G + (1-p)^G = mass of k in {0,G}  (p=0.05)|0.000e+00|tol=1.0e-15|PASS\n",
+      "          0.20    0.0281     0.150000       0.150000        6.713      2.3462 vs  2.3462\n",
+      "CHECK|exact|E[g_centred] = (1-1/G) p(1-p)  (p=0.2)|1.110e-16|tol=1.0e-15|PASS\n",
+      "CHECK|exact|E[sqrt(p_hat(1-p_hat))] / (p(1-p)) == c_G(p) of note 0105  (p=0.2)|4.441e-16|tol=1.0e-12|PASS\n",
+      "CHECK|exact|P(dead) = p^G + (1-p)^G = mass of k in {0,G}  (p=0.2)|0.000e+00|tol=1.0e-15|PASS\n",
+      "          0.50    0.0000     0.234375       0.234375      120.000      1.9341 vs  1.9341\n",
+      "CHECK|exact|E[g_centred] = (1-1/G) p(1-p)  (p=0.5)|0.000e+00|tol=1.0e-15|PASS\n",
+      "CHECK|exact|E[sqrt(p_hat(1-p_hat))] / (p(1-p)) == c_G(p) of note 0105  (p=0.5)|0.000e+00|tol=1.0e-12|PASS\n",
+      "CHECK|exact|P(dead) = p^G + (1-p)^G = mass of k in {0,G}  (p=0.5)|0.000e+00|tol=1.0e-15|PASS\n",
+      "          0.80    0.0281     0.150000       0.150000        6.713      2.3462 vs  2.3462\n",
+      "CHECK|exact|E[g_centred] = (1-1/G) p(1-p)  (p=0.8)|0.000e+00|tol=1.0e-15|PASS\n",
+      "CHECK|exact|E[sqrt(p_hat(1-p_hat))] / (p(1-p)) == c_G(p) of note 0105  (p=0.8)|0.000e+00|tol=1.0e-12|PASS\n",
+      "CHECK|exact|P(dead) = p^G + (1-p)^G = mass of k in {0,G}  (p=0.8)|0.000e+00|tol=1.0e-15|PASS\n",
+      "          0.95    0.4401     0.044531       0.044531        0.931      3.2671 vs  3.2671\n",
+      "CHECK|exact|E[g_centred] = (1-1/G) p(1-p)  (p=0.95)|6.939e-18|tol=1.0e-15|PASS\n",
+      "CHECK|exact|E[sqrt(p_hat(1-p_hat))] / (p(1-p)) == c_G(p) of note 0105  (p=0.95)|0.000e+00|tol=1.0e-12|PASS\n",
+      "CHECK|exact|P(dead) = p^G + (1-p)^G = mass of k in {0,G}  (p=0.95)|0.000e+00|tol=1.0e-15|PASS\n",
+      "          0.99    0.8515     0.009281       0.009281        0.165      3.7330 vs  3.7330\n",
+      "CHECK|exact|E[g_centred] = (1-1/G) p(1-p)  (p=0.99)|3.469e-18|tol=1.0e-15|PASS\n",
+      "CHECK|exact|E[sqrt(p_hat(1-p_hat))] / (p(1-p)) == c_G(p) of note 0105  (p=0.99)|4.441e-16|tol=1.0e-12|PASS\n",
+      "CHECK|exact|P(dead) = p^G + (1-p)^G = mass of k in {0,G}  (p=0.99)|0.000e+00|tol=1.0e-15|PASS\n",
+      "CHECK|exact|SNR of the group estimator peaks at p = 1/2 and decays towards both ends (monotone on each side)|0.000e+00|tol=0.0e+00|PASS\n",
+      "   ||group gradient|| for dead groups (4 samples):\n",
+      "     centre     all-correct: 0.0000    all-wrong: 0.0000\n",
+      "     std        all-correct: 0.0000    all-wrong: 0.0000\n",
+      "     rloo       all-correct: 0.0000    all-wrong: 0.0000\n",
+      "     plain01    all-correct: 2.1004    all-wrong: 0.0000\n",
+      "     plainpm    all-correct: 2.1004    all-wrong: 3.8422\n",
+      "CHECK|exact|dead groups have exactly zero gradient (centre)|0.000e+00|tol=0.0e+00|PASS\n",
+      "CHECK|exact|dead groups have exactly zero gradient (std)|0.000e+00|tol=0.0e+00|PASS\n",
+      "CHECK|exact|dead groups have exactly zero gradient (rloo)|0.000e+00|tol=0.0e+00|PASS\n",
+      "CHECK|exact|without a baseline an all-correct (R in {0,1}) or all-wrong (R in {-1,1}) group still moves the policy|0.000e+00|tol=0.0e+00|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Reward scale / sign conventions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "source": [
+    "gen = torch.Generator().manual_seed(1)\n",
+    "idx = torch.multinomial(p_seq, 8, replacement=True, generator=gen)\n",
+    "for kind in ('std',):\n",
+    "    g01 = group_gradient(idx, kind)\n",
+    "    Rpm_idx = idx\n",
+    "    Rg = 2 * R[idx] - 1\n",
+    "    sd = Rg.std(unbiased=False)\n",
+    "    gpm = (((Rg - Rg.mean()) / sd)[:, None] * S[idx]).mean(0) if float(sd) > 0 else torch.zeros_like(g01)\n",
+    "    check(\"GRPO (std-normalised): {0,1} and {-1,+1} rewards give the SAME gradient\", maxabs(g01, gpm), 1e-13)\n",
+    "g01 = group_gradient(idx, 'centre')\n",
+    "Rg = 2 * R[idx] - 1\n",
+    "gpm = ((Rg - Rg.mean())[:, None] * S[idx]).mean(0)\n",
+    "check(\"Dr.GRPO (centred): {-1,+1} rewards give exactly 2x the {0,1} gradient\", maxabs(gpm, 2 * g01), 1e-13)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|GRPO (std-normalised): {0,1} and {-1,+1} rewards give the SAME gradient|0.000e+00|tol=1.0e-13|PASS\n",
+      "CHECK|exact|Dr.GRPO (centred): {-1,+1} rewards give exactly 2x the {0,1} gradient|0.000e+00|tol=1.0e-13|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Potential-based shaping (finite horizon, gamma = 1):  r_t = Phi(s_{t+1}) - Phi(s_t),  Phi(terminal) = 0   => shaped return = R - Phi(s_0)  => same gradient"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "source": [
+    "gen = torch.Generator().manual_seed(3)\n",
+    "Phi = {}\n",
+    "for pre in lm.prefixes(): Phi[pre] = float(torch.randn(1, generator=gen)) * 2.0\n",
+    "for sq in lm.seqs:\n",
+    "    Phi.setdefault(sq, 0.0)                              # terminal states: Phi = 0\n",
+    "shaping = torch.zeros(lm.N)\n",
+    "for i, sq in enumerate(lm.seqs):\n",
+    "    tot = 0.0\n",
+    "    for t in range(len(sq)):\n",
+    "        a, b = sq[:t], sq[:t + 1]\n",
+    "        tot += Phi[b] - Phi[a]\n",
+    "    shaping[i] = tot\n",
+    "check(\"telescoping: sum_t [Phi(s_t+1) - Phi(s_t)] = -Phi(s_0) for every sequence (Phi(terminal)=0)\", float((shaping + Phi[()]).abs().max()), 1e-14)\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "g_shaped = torch.autograd.grad((lm.probs(th) * (R + shaping)).sum(), th)[0]\n",
+    "check(\"potential-based shaping leaves grad J unchanged (exact)\", maxabs(g_shaped, grad_p), 1e-14)\n",
+    "# a non-potential 'shaping' (terminal potential != 0, e.g. depending on length) changes the objective\n",
+    "Phi_bad = dict(Phi)\n",
+    "for sq in lm.seqs: Phi_bad[sq] = 0.7 * len(sq)\n",
+    "bad = torch.zeros(lm.N)\n",
+    "for i, sq in enumerate(lm.seqs):\n",
+    "    bad[i] = sum(Phi_bad[sq[:t + 1]] - Phi_bad[sq[:t]] for t in range(len(sq)))\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "g_bad = torch.autograd.grad((lm.probs(th) * (R + bad)).sum(), th)[0]\n",
+    "print(f\"   terminal potential 0.7*|o|: ||grad - grad J|| = {float((g_bad-grad_p).norm()):.4f}   (||grad J|| = {float(grad_p.norm()):.4f})\")\n",
+    "assert float((g_bad - grad_p).norm()) > 1e-3\n",
+    "check(\"a terminal potential (length-dependent) is NOT neutral: the gradient changes\", 0.0, 0.0)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|telescoping: sum_t [Phi(s_t+1) - Phi(s_t)] = -Phi(s_0) for every sequence (Phi(terminal)=0)|4.441e-16|tol=1.0e-14|PASS\n",
+      "CHECK|exact|potential-based shaping leaves grad J unchanged (exact)|1.457e-16|tol=1.0e-14|PASS\n",
+      "   terminal potential 0.7*|o|: ||grad - grad J|| = 0.5800   (||grad J|| = 0.1825)\n",
+      "CHECK|exact|a terminal potential (length-dependent) is NOT neutral: the gradient changes|0.000e+00|tol=0.0e+00|PASS\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Length penalties: as REWARD the objective is separable;  length has no other route to the policy"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "source": [
+    "lens = lm.lengths\n",
+    "lam = 0.05\n",
+    "b = float(p_seq @ R)\n",
+    "g_rew = (p_seq[:, None] * (R - lam * lens - (b - lam * float(p_seq @ lens)))[:, None] * S).sum(0)      # baselined REINFORCE with shaped reward\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "gJ = torch.autograd.grad((lm.probs(th) * R).sum(), th)[0]\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "gL = torch.autograd.grad((lm.probs(th) * lens).sum(), th)[0]\n",
+    "check(\"reward R - lam|o| : gradient = grad E[R] - lam grad E|o|  (separable, exact)\", maxabs(g_rew, gJ - lam * gL), 1e-14)\n",
+    "check(\"grad E|o| = E[|o| score]  (length can only reach the policy through the score function)\", maxabs(gL, (p_seq[:, None] * lens[:, None] * S).sum(0)), 1e-14)\n",
+    "# effective per-sequence rewards of the implicit length weighting of note 0105 vs an explicit penalty\n",
+    "R_eff_len = (R - b) / lens\n",
+    "print(\"   implicit reward of the 1/|o|-normalised GRPO loss:  (R-b)/|o|  -> not separable, depends on b and multiplies the reward by 1/|o|;\")\n",
+    "print(\"   explicit penalty R - lam|o| is separable and its trade-off is the single number lam.\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|reward R - lam|o| : gradient = grad E[R] - lam grad E|o|  (separable, exact)|1.041e-17|tol=1.0e-14|PASS\n",
+      "CHECK|exact|grad E|o| = E[|o| score]  (length can only reach the policy through the score function)|1.110e-16|tol=1.0e-14|PASS\n",
+      "   implicit reward of the 1/|o|-normalised GRPO loss:  (R-b)/|o|  -> not separable, depends on b and multiplies the reward by 1/|o|;\n",
+      "   explicit penalty R - lam|o| is separable and its trade-off is the single number lam.\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Verifier errors.  Constant false-accept / false-reject rates only rescale the gradient;  an exploitable false-accept class re-targets it."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "source": [
+    "aFA, aFR = 0.10, 0.05\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "Pc = (lm.probs(th) * R).sum()\n",
+    "Jobs = aFA + (1 - aFR - aFA) * Pc                                  # E[R_obs] when errors are independent of the response\n",
+    "Robs_expect = R * (1 - aFR) + (1 - R) * aFA                        # per-response expected observed reward\n",
+    "th2 = theta.clone().requires_grad_(True)\n",
+    "Jobs2 = (lm.probs(th2) * Robs_expect).sum()\n",
+    "g1 = torch.autograd.grad(Jobs, th)[0]; g2 = torch.autograd.grad(Jobs2, th2)[0]\n",
+    "check(\"E[R_obs] = aFA + (1-aFR-aFA) P(correct)  (value)\", abs(float(Jobs.detach()) - float(Jobs2.detach())), 1e-15)\n",
+    "check(\"constant error rates: grad E[R_obs] = (1-aFR-aFA) grad P(correct)  (exact rescaling)\", maxabs(g2, (1 - aFR - aFA) * grad_p), 1e-14)\n",
+    "# exploitable class: short sequences (length 2 ending in EOS) with an even token get accepted with prob 1\n",
+    "expl = torch.tensor([1.0 if (len(sq) == 2 and sq[-1] == lm.EOS and sq[0] % 2 == 0) else 0.0 for sq in lm.seqs]) * (1 - R)\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "gE = torch.autograd.grad((lm.probs(th) * expl).sum(), th)[0]\n",
+    "P_E = float((p_seq * expl).sum())\n",
+    "print(f\"   exploit class (always accepted although wrong): P = {P_E:.4f};  ||grad P(correct)|| = {float(grad_p.norm()):.4f};  ||grad P(exploit)|| = {float(gE.norm()):.4f};  cosine = {float(grad_p@gE/(grad_p.norm()*gE.norm())):+.3f}\")\n",
+    "th = theta.clone().requires_grad_(True)\n",
+    "g_obs = torch.autograd.grad((lm.probs(th) * (R + expl)).sum(), th)[0]      # E[R_obs], R_obs = 1 on correct OR exploit\n",
+    "check(\"observed-reward gradient (autograd) == grad P(correct) + grad P(exploit)\", maxabs(g_obs, grad_p + gE), 1e-14)\n",
+    "print(\"   exploit share of the observed-reward gradient norm: %.2f\" % float(gE.norm() / g_obs.norm()))"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "CHECK|exact|E[R_obs] = aFA + (1-aFR-aFA) P(correct)  (value)|0.000e+00|tol=1.0e-15|PASS\n",
+      "CHECK|exact|constant error rates: grad E[R_obs] = (1-aFR-aFA) grad P(correct)  (exact rescaling)|2.429e-17|tol=1.0e-14|PASS\n",
+      "   exploit class (always accepted although wrong): P = 0.2462;  ||grad P(correct)|| = 0.1825;  ||grad P(exploit)|| = 0.2786;  cosine = +0.083\n",
+      "CHECK|exact|observed-reward gradient (autograd) == grad P(correct) + grad P(exploit)|2.776e-17|tol=1.0e-14|PASS\n",
+      "   exploit share of the observed-reward gradient norm: 0.81\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7. Group-size / pass@k arithmetic: how many samples until a group is live; the cost of rare successes"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "source": [
+    "print(\"   p        P(live), G=8    G=16     G=64     expected prompts*samples per live group (G=16)\")\n",
+    "for pp in (0.001, 0.01, 0.05, 0.2, 0.5):\n",
+    "    live = [1 - pp ** G - (1 - pp) ** G for G in (8, 16, 64)]\n",
+    "    print(f\"   {pp:5.3f}     {live[0]:8.4f}   {live[1]:8.4f}   {live[2]:8.4f}      {16/live[1]:10.1f}\")\n",
+    "check(\"P(live) = 1 - p^G - (1-p)^G = pass@G - p^G  (identity)\", abs((1 - 0.2 ** 16 - 0.8 ** 16) - ((1 - 0.8 ** 16) - 0.2 ** 16)), 1e-15)\n",
+    "print(\"done; checks:\", len(CHECKS))"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "   p        P(live), G=8    G=16     G=64     expected prompts*samples per live group (G=16)\n",
+      "   0.001       0.0080     0.0159     0.0620          1007.5\n",
+      "   0.010       0.0773     0.1485     0.4744           107.7\n",
+      "   0.050       0.3366     0.5599     0.9625            28.6\n",
+      "   0.200       0.8322     0.9719     1.0000            16.5\n",
+      "   0.500       0.9922     1.0000     1.0000            16.0\n",
+      "CHECK|exact|P(live) = 1 - p^G - (1-p)^G = pass@G - p^G  (identity)|0.000e+00|tol=1.0e-15|PASS\n",
+      "done; checks: 40\n"
+     ]
+    }
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

code/notebooks/0201-verifier-binary-reward-mdp.ipynb

@@ -0,0 +1,641 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "b4fbe81d",
+   "metadata": {},
+   "source": [
+    "# 0201 -- Verifier-in-the-loop as a binary-reward MDP\n",
+    "\n",
+    "Scratch-pad for note `0201`. Everything is fp64, CPU only (`CUDA_VISIBLE_DEVICES=` is set\n",
+    "before torch is imported). Each section states the identity, computes it the *naive* way\n",
+    "(full enumeration, explicit hand-derived formula), then the *efficient* way (autograd on a\n",
+    "log-sum-exp-stable policy, Monte Carlo), and asserts that they agree.\n",
+    "\n",
+    "Identities verified:\n",
+    "1. Score-function gradient `grad J = E[r grad log pi]` against autograd of the exact `J`.\n",
+    "2. Token decomposition `grad log pi(y|x) = sum_t grad log pi(y_t|s_t)` and the hand formula `onehot - softmax`.\n",
+    "3. Baseline identity `E[grad log pi] = 0`, so any `b(x)` leaves the gradient unbiased.\n",
+    "4. Exact variance of the one-sample estimator and the optimal constant baseline `b*`.\n",
+    "5. Monte Carlo estimator converges at the `1/sqrt(N)` rate.\n",
+    "6. Verifier holes: `E[G] = E[V] - pi(H)` and the softmax-PG ODE `dz_k/dt = pi_k (r_k - J)`.\n",
+    "7. Proxy reward model (noisy) vs exact verifier: best-of-n selection error and the `sigma sqrt(2 ln n)` winner's curse."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "3b78deed",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:49:39.638799Z",
+     "iopub.status.busy": "2026-10-14T08:49:39.638648Z",
+     "iopub.status.idle": "2026-10-14T08:49:41.790359Z",
+     "shell.execute_reply": "2026-10-14T08:49:41.789459Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n",
+    "import itertools, math, json\n",
+    "import numpy as np, torch\n",
+    "torch.set_default_dtype(torch.float64)\n",
+    "torch.manual_seed(0); np.random.seed(0)\n",
+    "CHECKS = []\n",
+    "def check(name, err, tol):\n",
+    "    CHECKS.append((name, float(err), float(tol)))\n",
+    "    assert err <= tol, (name, float(err), tol)\n",
+    "    print(f'{name:70s} err={float(err):.3e}  (tol {tol:.0e})')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0d637e40",
+   "metadata": {},
+   "source": [
+    "## 1. A tabular autoregressive policy over proofs\n",
+    "Vocabulary `{0,1,2}`, proof length `T=3` => 27 \"proofs\". The policy has one logit vector per prefix\n",
+    "(13 prefixes of length 0,1,2). The verifier `V(y)` is a fixed random 0/1 table with about 30% ones."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "cf164f87",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:49:41.792633Z",
+     "iopub.status.busy": "2026-10-14T08:49:41.792359Z",
+     "iopub.status.idle": "2026-10-14T08:49:41.884373Z",
+     "shell.execute_reply": "2026-10-14T08:49:41.883429Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "num prefixes 13 num sequences 27 p = E[r] at init will be printed below\n",
+      "J(theta0) = 0.30454236036997673\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/vignesh/.opencode/tmp/ipykernel_442065/734446241.py:26: UserWarning: Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.\n",
+      "Consider using tensor.detach() first. (Triggered internally at /pytorch/torch/csrc/autograd/generated/python_variable_methods.cpp:838.)\n",
+      "  print('J(theta0) =', float(J))\n"
+     ]
+    }
+   ],
+   "source": [
+    "Vsz, T = 3, 3\n",
+    "prefixes = [()]\n",
+    "for L in range(1, T):\n",
+    "    prefixes += list(itertools.product(range(Vsz), repeat=L))\n",
+    "pid = {p: i for i, p in enumerate(prefixes)}\n",
+    "seqs = list(itertools.product(range(Vsz), repeat=T))\n",
+    "rng = np.random.RandomState(3)\n",
+    "r_tab = torch.tensor((rng.rand(len(seqs)) < 0.3).astype(float))   # V(y) in {0,1}\n",
+    "print('num prefixes', len(prefixes), 'num sequences', len(seqs), 'p = E[r] at init will be printed below')\n",
+    "\n",
+    "theta0 = torch.randn(len(prefixes), Vsz) * 0.7\n",
+    "\n",
+    "def log_pi_seq(theta, y):\n",
+    "    \"\"\"log pi(y) = sum_t log softmax(theta[prefix_t])[y_t]  (log_softmax is the stable form).\"\"\"\n",
+    "    lp = 0.0\n",
+    "    for t in range(T):\n",
+    "        lp = lp + torch.log_softmax(theta[pid[y[:t]]], dim=-1)[y[t]]\n",
+    "    return lp\n",
+    "\n",
+    "def exact_J(theta):\n",
+    "    return sum(torch.exp(log_pi_seq(theta, y)) * r_tab[i] for i, y in enumerate(seqs))\n",
+    "\n",
+    "th = theta0.clone().requires_grad_(True)\n",
+    "J = exact_J(th); J.backward()\n",
+    "g_auto = th.grad.clone()\n",
+    "print('J(theta0) =', float(J))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "78aebc59",
+   "metadata": {},
+   "source": [
+    "## 2. NAIVE hand-derived gradient: sum over all 27 sequences\n",
+    "`grad J = sum_y pi(y) r(y) sum_t grad log pi(y_t | s_t)`, with\n",
+    "`grad_{theta[s,:]} log pi(y_t|s_t) = onehot(y_t) - softmax(theta[s,:])` (the softmax-score identity, derived in the note)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "c3545a51",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:49:41.885962Z",
+     "iopub.status.busy": "2026-10-14T08:49:41.885798Z",
+     "iopub.status.idle": "2026-10-14T08:49:41.895143Z",
+     "shell.execute_reply": "2026-10-14T08:49:41.894090Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1. score-function grad (hand, enumerated) vs autograd of exact J       err=1.041e-17  (tol 1e-14)\n"
+     ]
+    }
+   ],
+   "source": [
+    "def hand_grad(theta):\n",
+    "    th_ = theta.detach()\n",
+    "    g = torch.zeros_like(th_)\n",
+    "    for i, y in enumerate(seqs):\n",
+    "        py = float(torch.exp(log_pi_seq(th_, y)))\n",
+    "        for t in range(T):\n",
+    "            s = pid[y[:t]]\n",
+    "            sm = torch.softmax(th_[s], dim=-1)\n",
+    "            oh = torch.zeros(Vsz); oh[y[t]] = 1.0\n",
+    "            g[s] += py * r_tab[i] * (oh - sm)\n",
+    "    return g\n",
+    "\n",
+    "g_hand = hand_grad(theta0)\n",
+    "check('1. score-function grad (hand, enumerated) vs autograd of exact J', (g_hand - g_auto).abs().max(), 1e-14)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "da5af085",
+   "metadata": {},
+   "source": [
+    "## 3. Baseline identity: `E[grad log pi] = 0`\n",
+    "Because `sum_y pi(y) = 1` for every theta, differentiate it: `0 = sum_y grad pi(y) = sum_y pi(y) grad log pi(y)`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "097f8049",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:49:41.896729Z",
+     "iopub.status.busy": "2026-10-14T08:49:41.896575Z",
+     "iopub.status.idle": "2026-10-14T08:49:41.926284Z",
+     "shell.execute_reply": "2026-10-14T08:49:41.924902Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "3. E[grad log pi] = 0 (baseline identity)                              err=5.898e-17  (tol 1e-15)\n",
+      "3. grad unchanged for b in {-3,0,0.3,17}                               err=6.557e-16  (tol 1e-13)\n"
+     ]
+    }
+   ],
+   "source": [
+    "def expected_score(theta):\n",
+    "    th_ = theta.detach(); g = torch.zeros_like(th_)\n",
+    "    for i, y in enumerate(seqs):\n",
+    "        py = float(torch.exp(log_pi_seq(th_, y)))\n",
+    "        for t in range(T):\n",
+    "            s = pid[y[:t]]; sm = torch.softmax(th_[s], dim=-1)\n",
+    "            oh = torch.zeros(Vsz); oh[y[t]] = 1.0\n",
+    "            g[s] += py * (oh - sm)\n",
+    "    return g\n",
+    "check('3. E[grad log pi] = 0 (baseline identity)', expected_score(theta0).abs().max(), 1e-15)\n",
+    "\n",
+    "# Baseline b constant in y: the gradient E[(r-b) grad log pi] is the same for any b.\n",
+    "def hand_grad_b(theta, b):\n",
+    "    th_ = theta.detach(); g = torch.zeros_like(th_)\n",
+    "    for i, y in enumerate(seqs):\n",
+    "        py = float(torch.exp(log_pi_seq(th_, y)))\n",
+    "        for t in range(T):\n",
+    "            s = pid[y[:t]]; sm = torch.softmax(th_[s], dim=-1)\n",
+    "            oh = torch.zeros(Vsz); oh[y[t]] = 1.0\n",
+    "            g[s] += py * (r_tab[i] - b) * (oh - sm)\n",
+    "    return g\n",
+    "errs = [(hand_grad_b(theta0, b) - g_auto).abs().max() for b in (-3.0, 0.0, 0.3, 17.0)]\n",
+    "check('3. grad unchanged for b in {-3,0,0.3,17}', max(errs), 1e-13)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "30ffac82",
+   "metadata": {},
+   "source": [
+    "## 4. Exact variance of the one-sample estimator and the optimal constant baseline\n",
+    "`Var[(r-b) s] = E[(r-b)^2 |s|^2] - |grad J|^2` with `s = grad log pi(y)` flattened. Minimiser:\n",
+    "`b* = E[r |s|^2] / E[|s|^2]`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "944489cb",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:49:41.928225Z",
+     "iopub.status.busy": "2026-10-14T08:49:41.927849Z",
+     "iopub.status.idle": "2026-10-14T08:49:43.333887Z",
+     "shell.execute_reply": "2026-10-14T08:49:43.332721Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "p = E[r] = 0.3045423603699767  b* = 0.33818921921147005  Var(b=0) = 0.5990480197782424  Var(b=p) = 0.39409891194469415  Var(b*) = 0.3920499407669431\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "4. argmin_b Var (grid) vs closed-form b*                               err=1.078e-05  (tol 2e-04)\n",
+      "4. Var(b*) closed form vs direct                                       err=5.551e-17  (tol 1e-12)\n",
+      "4. MC variance of one-sample estimator vs exact (rel. err.)            err=2.500e-03  (tol 2e-02)\n"
+     ]
+    }
+   ],
+   "source": [
+    "def score_vec(theta, y):\n",
+    "    th_ = theta.detach(); v = torch.zeros_like(th_)\n",
+    "    for t in range(T):\n",
+    "        s = pid[y[:t]]; sm = torch.softmax(th_[s], dim=-1)\n",
+    "        oh = torch.zeros(Vsz); oh[y[t]] = 1.0\n",
+    "        v[s] += oh - sm\n",
+    "    return v.flatten()\n",
+    "\n",
+    "P = torch.stack([torch.exp(log_pi_seq(theta0, y)).detach() for y in seqs])\n",
+    "S = torch.stack([score_vec(theta0, y) for y in seqs])          # (27, 39)\n",
+    "R = r_tab\n",
+    "gJ = (P[:, None] * R[:, None] * S).sum(0)\n",
+    "def var_exact(b):\n",
+    "    w = (R - b) ** 2 * (S ** 2).sum(1)\n",
+    "    return float((P * w).sum() - (gJ ** 2).sum())\n",
+    "bstar = float((P * R * (S ** 2).sum(1)).sum() / (P * (S ** 2).sum(1)).sum())\n",
+    "pbar = float((P * R).sum())\n",
+    "print('p = E[r] =', pbar, ' b* =', bstar, ' Var(b=0) =', var_exact(0.0), ' Var(b=p) =', var_exact(pbar), ' Var(b*) =', var_exact(bstar))\n",
+    "bs = np.linspace(-1, 2, 30001)\n",
+    "vs = np.array([var_exact(b) for b in bs])\n",
+    "check('4. argmin_b Var (grid) vs closed-form b*', abs(bs[vs.argmin()] - bstar), 2e-4)\n",
+    "# Closed form of the minimum: Var(b*) = E[r^2|s|^2] - (E[r|s|^2])^2/E|s|^2 - |gJ|^2  (r^2=r for binary r)\n",
+    "Es2 = float((P * (S ** 2).sum(1)).sum())\n",
+    "Ers2 = float((P * R * (S ** 2).sum(1)).sum())\n",
+    "var_min_closed = Ers2 - Ers2 ** 2 / Es2 - float((gJ ** 2).sum())\n",
+    "check('4. Var(b*) closed form vs direct', abs(var_min_closed - var_exact(bstar)), 1e-12)\n",
+    "# Monte Carlo of the one-sample estimator (independent check of var_exact at b = p)\n",
+    "N = 400000\n",
+    "idx = torch.multinomial(P, N, replacement=True)\n",
+    "est = (R[idx] - pbar)[:, None] * S[idx]\n",
+    "var_mc = float(((est - est.mean(0)) ** 2).sum(1).mean())\n",
+    "check('4. MC variance of one-sample estimator vs exact (rel. err.)', abs(var_mc - var_exact(pbar)) / var_exact(pbar), 0.02)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "db9b1afb",
+   "metadata": {},
+   "source": [
+    "## 5. Monte Carlo REINFORCE converges at the `1/sqrt(N)` rate\n",
+    "EFFICIENT version: sample `N` proofs from `pi`, average `(r - b) * score`. We compare with `grad J` in units of the\n",
+    "estimated standard error (z-scores), which must be O(1), and the RMS error must halve when `N` is quadrupled."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "64e6d7e6",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:49:43.335763Z",
+     "iopub.status.busy": "2026-10-14T08:49:43.335596Z",
+     "iopub.status.idle": "2026-10-14T08:49:44.334151Z",
+     "shell.execute_reply": "2026-10-14T08:49:44.332934Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "max |z| over 39 coordinates: 2.820402064782768\n",
+      "5. MC grad within 5 standard errors (max |z|, tol 5)                   err=2.820e+00  (tol 5e+00)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "RMS err N=5k: 0.008952773341135103  N=20k: 0.004564106878365233  ratio: 1.9615608441539822\n",
+      "5. rate: RMS(N)/RMS(4N) = 2 (abs dev)                                  err=3.844e-02  (tol 3e-01)\n"
+     ]
+    }
+   ],
+   "source": [
+    "def mc_grad(N, b, seed):\n",
+    "    gen = torch.Generator().manual_seed(seed)\n",
+    "    idx = torch.multinomial(P, N, replacement=True, generator=gen)\n",
+    "    e = (R[idx] - b)[:, None] * S[idx]\n",
+    "    return e.mean(0), e.std(0) / math.sqrt(N)\n",
+    "m, se = mc_grad(2_000_000, pbar, 11)\n",
+    "z = ((m - gJ) / se.clamp_min(1e-300))\n",
+    "mask = se > 1e-9\n",
+    "print('max |z| over', int(mask.sum()), 'coordinates:', float(z[mask].abs().max()))\n",
+    "check('5. MC grad within 5 standard errors (max |z|, tol 5)', float(z[mask].abs().max()), 5.0)\n",
+    "rms = lambda N: float(np.sqrt(np.mean([((mc_grad(N, pbar, s)[0] - gJ) ** 2).sum().item() for s in range(60)])))\n",
+    "r1, r4 = rms(5_000), rms(20_000)\n",
+    "print('RMS err N=5k:', r1, ' N=20k:', r4, ' ratio:', r1 / r4)\n",
+    "check('5. rate: RMS(N)/RMS(4N) = 2 (abs dev)', abs(r1 / r4 - 2.0), 0.35)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5a6f8e4c",
+   "metadata": {},
+   "source": [
+    "## 6. The hole: `E[G] = E[V] - pi(H)` and softmax policy-gradient dynamics\n",
+    "Three kinds of outputs: `wrong` (V=0,G=0), `proof` (V=1,G=1), `hole` (V=1, G=0, e.g. a `sorry`-like exploit that the\n",
+    "verifier accepts). For a tabular softmax over these classes, `dz_k/dt = pi_k (r_k - J)` (derived in the note)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "a8b49842",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:49:44.336095Z",
+     "iopub.status.busy": "2026-10-14T08:49:44.335921Z",
+     "iopub.status.idle": "2026-10-14T08:49:45.190504Z",
+     "shell.execute_reply": "2026-10-14T08:49:45.189496Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "6. softmax-PG gradient pi_k (r_k - J) vs autograd                      err=0.000e+00  (tol 1e-16)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "step, E[V], E[G], pi(hole)\n",
+      "  (0, 0.04, 0.01, 0.03)\n",
+      "  (4000, 0.9997460159631938, 0.0017478329594184914, 0.9979981830037753)\n",
+      "  (8000, 0.9998739033985322, 0.0012328930126125394, 0.9986410103859197)\n",
+      "  (12000, 0.9999161452892915, 0.001005875916774838, 0.9989102693725167)\n",
+      "  (16000, 0.9999371907590497, 0.0008707924515763814, 0.9990663983074733)\n",
+      "  (20000, 0.9999497931524061, 0.0007786947944058345, 0.9991710983580002)\n",
+      "6. identity E[G] = E[V] - pi(H) at the end of training                 err=1.735e-17  (tol 1e-15)\n",
+      "final: E[V]=0.9999  E[G]=0.0008  pi(hole)=0.9992  (the true-proof share of the accepted mass fell from 0.25 to 0.001)\n",
+      "6. proof share of accepted mass is non-increasing (max positive increment) err=0.000e+00  (tol 1e-15)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "6. false-reject bias: E[V] - (1-delta)E[G] within 4 s.e.               err=3.385e-01  (tol 4e+00)\n"
+     ]
+    }
+   ],
+   "source": [
+    "r_V = torch.tensor([0.0, 1.0, 1.0]); r_G = torch.tensor([0.0, 1.0, 0.0])\n",
+    "z = torch.log(torch.tensor([0.96, 0.01, 0.03]))      # hole is 3x more reachable than the true proof\n",
+    "# (a) gradient of J_V wrt logits: autograd vs hand formula pi_k (r_k - J)\n",
+    "zz = z.clone().requires_grad_(True)\n",
+    "pi = torch.softmax(zz, 0); JV = (pi * r_V).sum(); JV.backward()\n",
+    "hand = pi.detach() * (r_V - JV.detach())\n",
+    "check('6. softmax-PG gradient pi_k (r_k - J) vs autograd', (zz.grad - hand).abs().max(), 1e-16)\n",
+    "# (b) integrate the gradient flow with small steps; track J_V, J_G, and pi(H)\n",
+    "zt = z.clone(); eta = 0.5; traj = []\n",
+    "for step in range(20001):\n",
+    "    pi = torch.softmax(zt, 0)\n",
+    "    JVt = float((pi * r_V).sum()); JGt = float((pi * r_G).sum())\n",
+    "    if step % 4000 == 0: traj.append((step, JVt, JGt, float(pi[2])))\n",
+    "    zt = zt + eta * pi * (r_V - JVt)\n",
+    "print('step, E[V], E[G], pi(hole)'); [print(' ', t) for t in traj]\n",
+    "pi = torch.softmax(zt, 0)\n",
+    "EV = float((pi * r_V).sum()); EG = float((pi * r_G).sum()); piH = float(pi[2])\n",
+    "check('6. identity E[G] = E[V] - pi(H) at the end of training', abs(EG - (EV - piH)), 1e-15)\n",
+    "assert EV > 0.99 and EG < 0.3, (EV, EG)\n",
+    "print(f'final: E[V]={EV:.4f}  E[G]={EG:.4f}  pi(hole)={piH:.4f}  (the true-proof share of the accepted mass fell from 0.25 to {EG/EV:.3f})')\n",
+    "# Closed form for the invariant: log(pi_proof/pi_hole) changes by sum (pi_proof - pi_hole)(1-J) eta; since pi_hole > pi_proof it\n",
+    "# keeps growing in favour of the hole. Check monotonicity of the proof share along the trajectory.\n",
+    "zt = z.clone(); shares = []\n",
+    "for step in range(3000):\n",
+    "    pi = torch.softmax(zt, 0); shares.append(float(pi[1] / (pi[1] + pi[2])))\n",
+    "    zt = zt + eta * pi * (r_V - (pi * r_V).sum())\n",
+    "check('6. proof share of accepted mass is non-increasing (max positive increment)', max(0.0, float(np.max(np.diff(shares)))), 1e-15)\n",
+    "\n",
+    "# (c) false rejects: V = G * (1 - timeout) with independent timeouts of prob delta => E[V] = (1-delta) E[G]\n",
+    "delta = 0.1; N = 4_000_000\n",
+    "gen = torch.Generator().manual_seed(5)\n",
+    "G = (torch.rand(N, generator=gen) < 0.37).double(); to = (torch.rand(N, generator=gen) < delta).double()\n",
+    "Vv = G * (1 - to)\n",
+    "se = float(Vv.std() / math.sqrt(N))\n",
+    "check('6. false-reject bias: E[V] - (1-delta)E[G] within 4 s.e.', abs(float(Vv.mean()) - (1 - delta) * 0.37) / se, 4.0)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5a8054aa",
+   "metadata": {},
+   "source": [
+    "## 7. Proxy reward model vs exact verifier at best-of-`n`\n",
+    "`n` candidates, each truly correct w.p. `p`. A proxy RM scores `s = G + sigma * xi`, `xi ~ N(0,1)`. Selecting the argmax of `s`\n",
+    "picks a correct candidate with probability\n",
+    "`P_RM(n) = int n p f_1(s) [p F_1(s) + (1-p) F_0(s)]^(n-1) ds`\n",
+    "(`f_1,F_1` the N(1,sigma^2) pdf/cdf, `f_0,F_0` the N(0,sigma^2) ones), while the exact verifier gets\n",
+    "`1-(1-p)^n`. The *score* the RM reports for its pick exceeds the true value by about `sigma sqrt(2 ln n)`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "7cea16ac",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:49:45.192317Z",
+     "iopub.status.busy": "2026-10-14T08:49:45.192135Z",
+     "iopub.status.idle": "2026-10-14T08:51:27.965442Z",
+     "shell.execute_reply": "2026-10-14T08:51:27.964619Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "n=  4 p=0.2 sigma=0.5: P_RM exact=0.51088  MC=0.51153  verifier(any correct)=0.59040\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "n= 16 p=0.1 sigma=1.0: P_RM exact=0.33706  MC=0.33767  verifier(any correct)=0.81470\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "n= 64 p=0.05 sigma=1.5: P_RM exact=0.17838  MC=0.17813  verifier(any correct)=0.96248\n",
+      "7. best-of-n pick-correct probability: quadrature vs MC (max z)        err=8.216e-01  (tol 4e+00)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "verifier minus RM at n=16,p=0.1 for sigma=0.25,0.5,1,2: ['0.0089', '0.2043', '0.4776', '0.6137']\n",
+      "n=      10  E[max] exact=1.5388   upper bound sqrt(2 ln n)=2.1460\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "n=     100  E[max] exact=2.5076   upper bound sqrt(2 ln n)=3.0349\n",
+      "n=   10000  E[max] exact=3.8516   upper bound sqrt(2 ln n)=4.2919\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "n= 1000000  E[max] exact=4.8629   upper bound sqrt(2 ln n)=5.2565\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "7. E[max of 100 N(0,1)]: quadrature vs MC (in standard errors, tol 4)  err=1.327e+00  (tol 4e+00)\n"
+     ]
+    }
+   ],
+   "source": [
+    "from math import erf, sqrt, exp, pi as PI\n",
+    "def Phi(x): return 0.5 * (1 + np.vectorize(erf)(x / sqrt(2)))\n",
+    "def p_rm(n, p, sigma, grid=400001):\n",
+    "    s = np.linspace(-12 * sigma - 2, 12 * sigma + 3, grid)\n",
+    "    f1 = np.exp(-0.5 * ((s - 1) / sigma) ** 2) / (sigma * sqrt(2 * PI))\n",
+    "    F1 = Phi((s - 1) / sigma); F0 = Phi(s / sigma)\n",
+    "    integrand = n * p * f1 * (p * F1 + (1 - p) * F0) ** (n - 1)\n",
+    "    return float(np.trapezoid(integrand, s))\n",
+    "def p_rm_mc(n, p, sigma, trials, seed):\n",
+    "    g = np.random.RandomState(seed)\n",
+    "    G = (g.rand(trials, n) < p).astype(float); sc = G + sigma * g.randn(trials, n)\n",
+    "    pick = sc.argmax(1)\n",
+    "    return G[np.arange(trials), pick].mean(), sc.max(1).mean(), G[np.arange(trials), pick].mean()\n",
+    "errs = []\n",
+    "for (n, p, sig) in [(4, 0.2, 0.5), (16, 0.1, 1.0), (64, 0.05, 1.5)]:\n",
+    "    exact = p_rm(n, p, sig); mc, _, _ = p_rm_mc(n, p, sig, 400000, 1)\n",
+    "    se = math.sqrt(exact * (1 - exact) / 400000)\n",
+    "    errs.append(abs(exact - mc) / se)\n",
+    "    print(f'n={n:3d} p={p} sigma={sig}: P_RM exact={exact:.5f}  MC={mc:.5f}  verifier(any correct)={1-(1-p)**n:.5f}')\n",
+    "check('7. best-of-n pick-correct probability: quadrature vs MC (max z)', max(errs), 4.0)\n",
+    "# Verifier is never worse and strictly better for sigma>0\n",
+    "gap = [1 - (1 - 0.1) ** 16 - p_rm(16, 0.1, s) for s in (0.25, 0.5, 1.0, 2.0)]\n",
+    "print('verifier minus RM at n=16,p=0.1 for sigma=0.25,0.5,1,2:', [f'{x:.4f}' for x in gap])\n",
+    "assert all(g > 0 for g in gap) and all(np.diff(gap) > 0)\n",
+    "# Winner's curse: E[max of n iid N(0,1)] = int n x phi(x) Phi(x)^(n-1) dx (exact quadrature); bound E[max] <= sqrt(2 ln n).\n",
+    "def emax_exact(n):\n",
+    "    x = np.linspace(-10, 10, 400001)\n",
+    "    phi = np.exp(-0.5 * x ** 2) / sqrt(2 * PI)\n",
+    "    return float(np.trapezoid(n * x * phi * Phi(x) ** (n - 1), x))\n",
+    "for n in (10, 100, 10_000, 1_000_000):\n",
+    "    print(f'n={n:8d}  E[max] exact={emax_exact(n):.4f}   upper bound sqrt(2 ln n)={sqrt(2*math.log(n)):.4f}')\n",
+    "    assert emax_exact(n) <= sqrt(2 * math.log(n))\n",
+    "mcs = np.random.RandomState(2).randn(40000, 100).max(1)\n",
+    "check('7. E[max of 100 N(0,1)]: quadrature vs MC (in standard errors, tol 4)', abs(emax_exact(100) - mcs.mean()) / (mcs.std() / sqrt(len(mcs))), 4.0)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f1a420d1",
+   "metadata": {},
+   "source": [
+    "## Verification table"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "3c054d00",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:51:27.967075Z",
+     "iopub.status.busy": "2026-10-14T08:51:27.966912Z",
+     "iopub.status.idle": "2026-10-14T08:51:27.970746Z",
+     "shell.execute_reply": "2026-10-14T08:51:27.970067Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "identity                                                                  max abs err\n",
+      "1. score-function grad (hand, enumerated) vs autograd of exact J            1.041e-17\n",
+      "3. E[grad log pi] = 0 (baseline identity)                                   5.898e-17\n",
+      "3. grad unchanged for b in {-3,0,0.3,17}                                    6.557e-16\n",
+      "4. argmin_b Var (grid) vs closed-form b*                                    1.078e-05\n",
+      "4. Var(b*) closed form vs direct                                            5.551e-17\n",
+      "4. MC variance of one-sample estimator vs exact (rel. err.)                 2.500e-03\n",
+      "5. MC grad within 5 standard errors (max |z|, tol 5)                        2.820e+00\n",
+      "5. rate: RMS(N)/RMS(4N) = 2 (abs dev)                                       3.844e-02\n",
+      "6. softmax-PG gradient pi_k (r_k - J) vs autograd                           0.000e+00\n",
+      "6. identity E[G] = E[V] - pi(H) at the end of training                      1.735e-17\n",
+      "6. proof share of accepted mass is non-increasing (max positive increment)    0.000e+00\n",
+      "6. false-reject bias: E[V] - (1-delta)E[G] within 4 s.e.                    3.385e-01\n",
+      "7. best-of-n pick-correct probability: quadrature vs MC (max z)             8.216e-01\n",
+      "7. E[max of 100 N(0,1)]: quadrature vs MC (in standard errors, tol 4)       1.327e+00\n",
+      "CHECKS_JSON:[[\"1. score-function grad (hand, enumerated) vs autograd of exact J\", 1.0408340855860843e-17, 1e-14], [\"3. E[grad log pi] = 0 (baseline identity)\", 5.898059818321144e-17, 1e-15], [\"3. grad unchanged for b in {-3,0,0.3,17}\", 6.557254739192331e-16, 1e-13], [\"4. argmin_b Var (grid) vs closed-form b*\", 1.078078853000708e-05, 0.0002], [\"4. Var(b*) closed form vs direct\", 5.551115123125783e-17, 1e-12], [\"4. MC variance of one-sample estimator vs exact (rel. err.)\", 0.0025004228417077347, 0.02], [\"5. MC grad within 5 standard errors (max |z|, tol 5)\", 2.820402064782768, 5.0], [\"5. rate: RMS(N)/RMS(4N) = 2 (abs dev)\", 0.03843915584601776, 0.35], [\"6. softmax-PG gradient pi_k (r_k - J) vs autograd\", 0.0, 1e-16], [\"6. identity E[G] = E[V] - pi(H) at the end of training\", 1.734723475976807e-17, 1e-15], [\"6. proof share of accepted mass is non-increasing (max positive increment)\", 0.0, 1e-15], [\"6. false-reject bias: E[V] - (1-delta)E[G] within 4 s.e.\", 0.33845555650118325, 4.0], [\"7. best-of-n pick-correct probability: quadrature vs MC (max z)\", 0.8216357277155372, 4.0], [\"7. E[max of 100 N(0,1)]: quadrature vs MC (in standard errors, tol 4)\", 1.327150492424622, 4.0]]\n"
+     ]
+    }
+   ],
+   "source": [
+    "print()\n",
+    "print(f'{\"identity\":72s} {\"max abs err\":>12s}')\n",
+    "for n_, e_, t_ in CHECKS: print(f'{n_:72s} {e_:12.3e}')\n",
+    "print('CHECKS_JSON:' + json.dumps(CHECKS))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

code/notebooks/0202-rejection-sampling-expert-iteration-em.ipynb

@@ -0,0 +1,555 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "9658b1ce",
+   "metadata": {},
+   "source": [
+    "# 0202 -- Rejection sampling / expert iteration as policy improvement (and as EM)\n",
+    "\n",
+    "Identities verified (fp64, CPU):\n",
+    "1. `grad_theta E_x E_{y~pi_t}[V log pi_theta]` at `theta=theta_t` equals the policy gradient `grad J(theta_t)`.\n",
+    "2. ELBO identity `L(theta,q) = log p_theta(x) - KL(q || q*_theta)` and tightness at `q=q*_theta` (so the E-step is exact).\n",
+    "3. Tabular M-step: `argmin KL(q_t || pi_theta)` gives `pi_{t+1}=q_t`, hence `p_{t+1}(x)=1` where `p_t(x)>0` and `0` where `p_t(x)=0` (coverage).\n",
+    "4. Shared-parameter model: with an *exact* M-step, `F(theta)=sum_x log p_theta(x)` is non-decreasing (EM monotonicity).\n",
+    "5. Sampling: `E[ sum_i V_i f(y_i)/c ; c>0 ] = (1-(1-p)^n) E_q[f]` (exact enumeration) -- the self-normalised filtered average is unbiased except for the missing event `c=0`.\n",
+    "6. Unnormalised (practical) SFT objective is *not* monotone per prompt under shared parameters.\n",
+    "7. Monte Carlo RFT gradient agrees with the exact `grad J` (RFT step = REINFORCE step)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "86d1f2a7",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:47:57.273949Z",
+     "iopub.status.busy": "2026-10-14T09:47:57.273796Z",
+     "iopub.status.idle": "2026-10-14T09:47:59.029734Z",
+     "shell.execute_reply": "2026-10-14T09:47:59.029103Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n",
+    "import itertools, math, json\n",
+    "import numpy as np, torch\n",
+    "torch.set_default_dtype(torch.float64)\n",
+    "torch.manual_seed(1); np.random.seed(1)\n",
+    "CHECKS = []\n",
+    "def check(name, err, tol):\n",
+    "    CHECKS.append((name, float(err), float(tol)))\n",
+    "    assert err <= tol, (name, float(err), tol)\n",
+    "    print(f'{name:78s} err={float(err):.3e}  (tol {tol:.0e})')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ee2a1501",
+   "metadata": {},
+   "source": [
+    "## 1. A small family of prompts with a tabular policy\n",
+    "`m=6` prompts, `K=10` candidate proofs per prompt. `V[x,y]` in {0,1}. Prompt 5 has *no* correct proof; prompt 4 has a correct\n",
+    "proof that the initial policy almost never produces (coverage problem)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "5af4e833",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:47:59.032313Z",
+     "iopub.status.busy": "2026-10-14T09:47:59.032078Z",
+     "iopub.status.idle": "2026-10-14T09:47:59.039690Z",
+     "shell.execute_reply": "2026-10-14T09:47:59.038725Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "p_t(x) = [0.26066, 0.36631, 0.17998, 0.16817, 0.00018, 0.0]\n"
+     ]
+    }
+   ],
+   "source": [
+    "m, K = 6, 10\n",
+    "rng = np.random.RandomState(7)\n",
+    "V = (rng.rand(m, K) < 0.25).astype(float)\n",
+    "V[5, :] = 0.0                       # unsolvable: V == 0 everywhere\n",
+    "V[4, :] = 0.0; V[4, 9] = 1.0        # exactly one correct proof\n",
+    "V[0, 0] = 1.0                        # make sure prompt 0 is solvable\n",
+    "V = torch.tensor(V)\n",
+    "logits0 = torch.randn(m, K) * 0.5\n",
+    "logits0[4, 9] = -6.0                 # pi(correct) ~ 2e-4 for prompt 4\n",
+    "def pi_tab(z): return torch.softmax(z, dim=-1)\n",
+    "def p_of(z): return (pi_tab(z) * V).sum(-1)\n",
+    "print('p_t(x) =', [round(float(v), 5) for v in p_of(logits0)])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7bf5412a",
+   "metadata": {},
+   "source": [
+    "## 2. First-order equivalence: one RFT step is a policy-gradient step\n",
+    "RFT objective with samples from the *frozen* `pi_t`: `L(theta) = (1/m) sum_x sum_y pi_t(y|x) V(x,y) log pi_theta(y|x)`.\n",
+    "Its gradient at `theta=theta_t` is `(1/m) sum_x sum_y V pi_t grad log pi_t = grad J(theta_t)` (score identity of note 0201)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "d29b7d32",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:47:59.041294Z",
+     "iopub.status.busy": "2026-10-14T09:47:59.041085Z",
+     "iopub.status.idle": "2026-10-14T09:47:59.128842Z",
+     "shell.execute_reply": "2026-10-14T09:47:59.127645Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1. grad of RFT objective at theta_t vs grad J(theta_t)                         err=3.469e-18  (tol 1e-16)\n"
+     ]
+    }
+   ],
+   "source": [
+    "z = logits0.clone().requires_grad_(True)\n",
+    "J = p_of(z).mean(); J.backward(); gJ = z.grad.clone()\n",
+    "pi_t = pi_tab(logits0).detach()\n",
+    "z2 = logits0.clone().requires_grad_(True)\n",
+    "L = ((pi_t * V * torch.log_softmax(z2, -1)).sum(-1)).mean(); L.backward()\n",
+    "check('1. grad of RFT objective at theta_t vs grad J(theta_t)', (z2.grad - gJ).abs().max(), 1e-16)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "93fcb188",
+   "metadata": {},
+   "source": [
+    "## 3. ELBO identity and exact E-step\n",
+    "`p_theta(x) = sum_y pi_theta V`. `q*_theta(y|x) = pi_theta V / p_theta(x)`. For any `q` (full support on {V=1}):\n",
+    "`L(theta,q) = E_q[log V(x,y)] - KL(q || pi_theta)`, and with `V in {0,1}` and `q` supported on `{V=1}` this is `-KL(q||pi_theta)`.\n",
+    "Claim: `-KL(q||pi_theta) = log p_theta - KL(q||q*_theta)`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "5dd4e7bb",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:47:59.130670Z",
+     "iopub.status.busy": "2026-10-14T09:47:59.130469Z",
+     "iopub.status.idle": "2026-10-14T09:47:59.137329Z",
+     "shell.execute_reply": "2026-10-14T09:47:59.136038Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2. ELBO identity  -KL(q||pi) = log p - KL(q||q*)                               err=0.000e+00  (tol 1e-14)\n",
+      "2. tightness: -KL(q*||pi) = log p                                              err=2.220e-16  (tol 1e-14)\n"
+     ]
+    }
+   ],
+   "source": [
+    "th = logits0[0]; pi = pi_tab(th); Vx = V[0]; p0 = float((pi * Vx).sum())\n",
+    "qstar = pi * Vx / p0\n",
+    "qrand = torch.rand(K) * Vx; qrand = qrand / qrand.sum()\n",
+    "def KL(a, b):\n",
+    "    msk = a > 0\n",
+    "    return float((a[msk] * (torch.log(a[msk]) - torch.log(b[msk]))).sum())\n",
+    "lhs = -KL(qrand, pi)\n",
+    "rhs = math.log(p0) - KL(qrand, qstar)\n",
+    "check('2. ELBO identity  -KL(q||pi) = log p - KL(q||q*)', abs(lhs - rhs), 1e-14)\n",
+    "check('2. tightness: -KL(q*||pi) = log p', abs(-KL(qstar, pi) - math.log(p0)), 1e-14)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ae2285be",
+   "metadata": {},
+   "source": [
+    "## 4. Tabular M-step: `pi_{t+1} = q_t` and the coverage theorem\n",
+    "With full capacity `argmin_theta KL(q_t||pi_theta)` gives `pi_{t+1}(.|x)=q_t(.|x)` for every prompt with `p_t(x)>0`.\n",
+    "One iteration therefore drives `p(x)` to 1 on all covered prompts and leaves `p(x)=0` prompts untouched."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "d07a7eb0",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:47:59.139284Z",
+     "iopub.status.busy": "2026-10-14T09:47:59.139078Z",
+     "iopub.status.idle": "2026-10-14T09:47:59.146022Z",
+     "shell.execute_reply": "2026-10-14T09:47:59.144837Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "p_t   : [0.26066, 0.36631, 0.17998, 0.16817, 0.00018, 0.0]\n",
+      "p_t+1 : [1.0, 1.0, 1.0, 1.0, 1.0, 0.0]\n",
+      "3. p_{t+1}(x) = 1 on covered prompts (max |p-1|)                               err=0.000e+00  (tol 1e-14)\n",
+      "3. p_{t+1}(x) = 0 on uncovered prompt (coverage failure)                       err=0.000e+00  (tol 0e+00)\n"
+     ]
+    }
+   ],
+   "source": [
+    "pt = p_of(logits0)\n",
+    "covered = pt > 0\n",
+    "q = pi_tab(logits0) * V\n",
+    "q = torch.where(covered[:, None], q / pt.clamp_min(1e-300)[:, None], pi_tab(logits0))   # keep pi_t where q undefined\n",
+    "pi_next = q\n",
+    "p_next = (pi_next * V).sum(-1)\n",
+    "print('p_t   :', [round(float(v), 5) for v in pt])\n",
+    "print('p_t+1 :', [round(float(v), 5) for v in p_next])\n",
+    "check('3. p_{t+1}(x) = 1 on covered prompts (max |p-1|)', (p_next[covered] - 1).abs().max(), 1e-14)\n",
+    "check('3. p_{t+1}(x) = 0 on uncovered prompt (coverage failure)', p_next[~covered].abs().max(), 0.0)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3670c24a",
+   "metadata": {},
+   "source": [
+    "## 5. Shared parameters: EM monotonicity with an exact M-step\n",
+    "Log-linear policy `pi_theta(y|x) = softmax_y(phi(x,y) . theta)`, `d=5 < K=12` features, so the model cannot represent `q_t` exactly.\n",
+    "`F(theta) = (1/|S|) sum_{x in S} log p_theta(x)` over the prompts with at least one correct proof (`S`).\n",
+    "The M-step objective `-(1/|S|) sum_x sum_y q_t(y|x) log pi_theta(y|x)` is convex in theta (softmax cross-entropy with linear logits),\n",
+    "so L-BFGS solves it to machine precision, i.e. the M-step is exact."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "341bebce",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:47:59.148100Z",
+     "iopub.status.busy": "2026-10-14T09:47:59.147866Z",
+     "iopub.status.idle": "2026-10-14T09:48:00.351047Z",
+     "shell.execute_reply": "2026-10-14T09:48:00.350146Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "F per iteration: [-1.733346, -1.418455, -1.386263, -1.381133, -1.377459, -1.373836, -1.37042, -1.367474, -1.365174, -1.363539, -1.362461, -1.361784, -1.361367]\n",
+      "mean p per iter: [0.185491, 0.234027, 0.253918, 0.259472, 0.261745, 0.263587, 0.265459, 0.267323, 0.269065, 0.270608, 0.271926, 0.273028, 0.273945]\n",
+      "final M-step max |grad|: 1.2179568404172003e-09\n",
+      "4. EM monotone: max decrease of F over 12 iterations (should be <= 0)          err=0.000e+00  (tol 1e-10)\n",
+      "4. M-step solved to optimality (max |grad| at the last M-step)                 err=1.218e-09  (tol 1e-08)\n",
+      "4. ELBO(theta_t,q_t) = F(theta_t) at the final iterate                         err=0.000e+00  (tol 1e-12)\n",
+      "4. equal-weight M-step objective: gradient at theta_t = grad F(theta_t) (weight 1/p_t per problem on grad p) err=6.245e-17  (tol 1e-14)\n",
+      "4. pooled objective (ReST-EM Eq.3 / RFT): gradient at theta_t = grad mean_S p_t (weight 1 on grad p) err=1.388e-17  (tol 1e-14)\n",
+      "relative difference of the two directions (cosine): 0.10628750941093855\n"
+     ]
+    }
+   ],
+   "source": [
+    "m2, K2, d = 8, 12, 5\n",
+    "g = np.random.RandomState(11)\n",
+    "Phi = torch.tensor(g.randn(m2, K2, d))\n",
+    "V2 = (g.rand(m2, K2) < 0.2).astype(float)\n",
+    "V2[7, :] = 0.0                                   # unsolvable prompt\n",
+    "for x in range(7):\n",
+    "    if V2[x].sum() == 0: V2[x, x] = 1.0\n",
+    "V2 = torch.tensor(V2)\n",
+    "S = (V2.sum(-1) > 0)\n",
+    "def logpi(theta): return torch.log_softmax(Phi @ theta, dim=-1)           # (m2, K2)\n",
+    "def p_x(theta): return (torch.exp(logpi(theta)) * V2).sum(-1)\n",
+    "def F(theta): return torch.log(p_x(theta)[S]).mean()\n",
+    "def e_step(theta):\n",
+    "    w = torch.exp(logpi(theta)) * V2\n",
+    "    return w / w.sum(-1, keepdim=True).clamp_min(1e-300)\n",
+    "def m_step(q, theta_init, weights=None):\n",
+    "    th_ = theta_init.clone().requires_grad_(True)\n",
+    "    wts = torch.ones(m2) if weights is None else weights\n",
+    "    opt = torch.optim.LBFGS([th_], lr=1.0, max_iter=500, tolerance_grad=1e-13, tolerance_change=1e-16, line_search_fn='strong_wolfe')\n",
+    "    def closure():\n",
+    "        opt.zero_grad()\n",
+    "        loss = -(wts[S, None] * q[S] * logpi(th_)[S]).sum() / S.sum()\n",
+    "        loss.backward(); return loss\n",
+    "    for _ in range(4): opt.step(closure)\n",
+    "    return th_.detach(), th_.grad.abs().max().item()\n",
+    "theta = torch.randn(d) * 0.1\n",
+    "Fs, Js, gmax = [float(F(theta))], [float(p_x(theta).mean())], []\n",
+    "for it in range(12):\n",
+    "    q = e_step(theta)\n",
+    "    theta, gm = m_step(q, theta); gmax.append(gm)\n",
+    "    Fs.append(float(F(theta))); Js.append(float(p_x(theta).mean()))\n",
+    "print('F per iteration:', [round(v, 6) for v in Fs])\n",
+    "print('mean p per iter:', [round(v, 6) for v in Js])\n",
+    "print('final M-step max |grad|:', gmax[-1])\n",
+    "check('4. EM monotone: max decrease of F over 12 iterations (should be <= 0)', max(0.0, max(-(b - a) for a, b in zip(Fs[:-1], Fs[1:]))), 1e-10)\n",
+    "check('4. M-step solved to optimality (max |grad| at the last M-step)', gmax[-1], 1e-8)\n",
+    "assert Fs[-1] > Fs[0] + 0.05\n",
+    "\n",
+    "# Tightness at theta_t, per iteration: ELBO(theta_t, q_t) = F(theta_t)\n",
+    "q = e_step(theta)\n",
+    "elbo = -(q[S] * (torch.log(q[S].clamp_min(1e-300)) - logpi(theta)[S]) * (V2[S] > 0)).sum(-1).mean()\n",
+    "check('4. ELBO(theta_t,q_t) = F(theta_t) at the final iterate', abs(float(elbo) - float(F(theta))), 1e-12)\n",
+    "\n",
+    "# Per-problem weights of the two M-step objectives, checked at theta_t (this is the 1 vs 1/p_t(x) claim of the comparison table):\n",
+    "#   equal-weight EM objective  (1/|S|) sum_x E_{q_t}[log pi_theta]  has gradient (1/|S|) sum_x grad log p_t(x) = grad F(theta_t)\n",
+    "#   pooled RFT / ReST-EM Eq.(3) objective (1/|S|) sum_x p_t(x) E_{q_t}[log pi_theta]  has gradient (1/|S|) sum_x grad p_t(x) = grad (mean_S p)\n",
+    "th_chk = theta.clone().requires_grad_(True)\n",
+    "qt = e_step(theta).detach()\n",
+    "(-(qt[S] * logpi(th_chk)[S]).sum() / S.sum()).backward(); g_eq = -th_chk.grad.clone()\n",
+    "th_f = theta.clone().requires_grad_(True); F(th_f).backward()\n",
+    "check('4. equal-weight M-step objective: gradient at theta_t = grad F(theta_t) (weight 1/p_t per problem on grad p)', (g_eq - th_f.grad).abs().max(), 1e-14)\n",
+    "th_chk = theta.clone().requires_grad_(True)\n",
+    "pw = p_x(theta).detach()\n",
+    "(-(pw[S, None] * qt[S] * logpi(th_chk)[S]).sum() / S.sum()).backward(); g_pool = -th_chk.grad.clone()\n",
+    "th_p = theta.clone().requires_grad_(True); p_x(th_p)[S].mean().backward()\n",
+    "check('4. pooled objective (ReST-EM Eq.3 / RFT): gradient at theta_t = grad mean_S p_t (weight 1 on grad p)', (g_pool - th_p.grad).abs().max(), 1e-14)\n",
+    "print('relative difference of the two directions (cosine):', float(g_eq @ g_pool / (g_eq.norm() * g_pool.norm())))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4233ec2d",
+   "metadata": {},
+   "source": [
+    "## 6. The practical objective (unnormalised, weights `p_t(x)`) is not monotone per prompt\n",
+    "SFT on the pooled filtered samples weights prompt `x` by (expected) count of correct samples `n p_t(x)`, i.e. objective\n",
+    "`sum_x p_t(x) E_{q_t}[log pi_theta]`. With a shared `theta` some prompts get *worse* even though the weighted objective improves."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "fb963f24",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:48:00.353128Z",
+     "iopub.status.busy": "2026-10-14T09:48:00.352833Z",
+     "iopub.status.idle": "2026-10-14T09:48:00.421027Z",
+     "shell.execute_reply": "2026-10-14T09:48:00.420159Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "iter, mean p on solvable prompts, #prompts whose p decreased, F\n",
+      "  (0, 0.2615781869672757, 2, -1.4525900599954458)\n",
+      "  (1, 0.305815927384656, 3, -1.4159837962389783)\n",
+      "  (2, 0.360538810599986, 4, -1.5299987749300417)\n",
+      "  (3, 0.3968721919898336, 4, -1.7930711977638583)\n",
+      "  (4, 0.4128722015161725, 4, -2.0991072756580036)\n",
+      "  (5, 0.42201137651015436, 3, -2.3539650242141725)\n",
+      "total (prompt, iteration) pairs with decreasing p: 20\n"
+     ]
+    }
+   ],
+   "source": [
+    "theta_u = torch.randn(d) * 0.1\n",
+    "rows = []\n",
+    "for it in range(6):\n",
+    "    p_before = p_x(theta_u).clone()\n",
+    "    q = e_step(theta_u)\n",
+    "    theta_u, gm = m_step(q, theta_u, weights=p_before.detach())\n",
+    "    p_after = p_x(theta_u)\n",
+    "    rows.append((it, float(p_after[S].mean()), int((p_after[S] < p_before[S] - 1e-12).sum()), float(F(theta_u))))\n",
+    "print('iter, mean p on solvable prompts, #prompts whose p decreased, F')\n",
+    "for r in rows: print(' ', r)\n",
+    "n_dec = sum(r[2] for r in rows)\n",
+    "print('total (prompt, iteration) pairs with decreasing p:', n_dec)\n",
+    "CHECKS.append(('6. non-monotone prompts exist under unnormalised weights (count of decreasing p, informational)', float(n_dec), float('inf')))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ba81c4f7",
+   "metadata": {},
+   "source": [
+    "## 7. Sampling estimator: unbiased except for the event `c = 0`\n",
+    "One prompt, outcomes `y in {0,1,2,3}`, `V=(0,1,1,0)`, `pi=(.4,.1,.2,.3)`, a test function `f`. Draw `n=3` i.i.d. samples; keep correct ones;\n",
+    "estimate `E_q[f]` by `(1/c) sum_i V_i f(y_i)` when `c>0` (else 0). Claim: expectation `= (1-(1-p)^n) E_q f` exactly.\n",
+    "NAIVE: enumerate all `4^3` sample tuples with their probabilities. EFFICIENT: the closed form."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "52c6844e",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:48:00.422626Z",
+     "iopub.status.busy": "2026-10-14T09:48:00.422467Z",
+     "iopub.status.idle": "2026-10-14T09:48:00.441049Z",
+     "shell.execute_reply": "2026-10-14T09:48:00.440191Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "5. E[filtered mean of f] = (1-(1-p)^n) E_q[f],  n=1                            err=0.000e+00  (tol 1e-14)\n",
+      "5. E[filtered mean of f] = (1-(1-p)^n) E_q[f],  n=2                            err=3.331e-16  (tol 1e-14)\n",
+      "5. E[filtered mean of f] = (1-(1-p)^n) E_q[f],  n=3                            err=8.882e-16  (tol 1e-14)\n",
+      "5. E[filtered mean of f] = (1-(1-p)^n) E_q[f],  n=5                            err=2.220e-15  (tol 1e-14)\n",
+      "5. pooled filtered average: E[(1/n) sum V f] = p E_q[f]                        err=3.331e-16  (tol 1e-14)\n"
+     ]
+    }
+   ],
+   "source": [
+    "pi4 = np.array([.4, .1, .2, .3]); V4 = np.array([0, 1, 1, 0.]); f4 = np.array([1.3, -0.7, 2.9, 5.0])\n",
+    "p = float((pi4 * V4).sum()); Eq_f = float((pi4 * V4 * f4).sum() / p)\n",
+    "for n in (1, 2, 3, 5):\n",
+    "    tot = 0.0\n",
+    "    for tup in itertools.product(range(4), repeat=n):\n",
+    "        prob = np.prod([pi4[i] for i in tup]); c = sum(V4[i] for i in tup)\n",
+    "        if c > 0: tot += prob * sum(V4[i] * f4[i] for i in tup) / c\n",
+    "    closed = (1 - (1 - p) ** n) * Eq_f\n",
+    "    check(f'5. E[filtered mean of f] = (1-(1-p)^n) E_q[f],  n={n}', abs(tot - closed), 1e-14)\n",
+    "\n",
+    "# Pooled (unnormalised) filtered average: E[(1/n) sum_i V_i f(y_i)] = p E_q[f] exactly (linearity).\n",
+    "n = 4; tot = 0.0\n",
+    "for tup in itertools.product(range(4), repeat=n):\n",
+    "    prob = np.prod([pi4[i] for i in tup]); tot += prob * sum(V4[i] * f4[i] for i in tup) / n\n",
+    "check('5. pooled filtered average: E[(1/n) sum V f] = p E_q[f]', abs(tot - p * Eq_f), 1e-14)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e2501a9f",
+   "metadata": {},
+   "source": [
+    "## 8. Monte Carlo agreement with the exact tabular gradient (RFT step = REINFORCE step)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "5aa7257d",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:48:00.442590Z",
+     "iopub.status.busy": "2026-10-14T09:48:00.442434Z",
+     "iopub.status.idle": "2026-10-14T09:48:00.525797Z",
+     "shell.execute_reply": "2026-10-14T09:48:00.524856Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "7. Monte Carlo RFT gradient (N=4e5 per prompt) vs grad J  (tol 3e-4)           err=1.941e-04  (tol 3e-04)\n"
+     ]
+    }
+   ],
+   "source": [
+    "N = 400_000\n",
+    "gen = torch.Generator().manual_seed(3)\n",
+    "Pt = pi_tab(logits0)\n",
+    "gmc = torch.zeros(m, K)\n",
+    "for x in range(m):\n",
+    "    ys = torch.multinomial(Pt[x], N, replacement=True, generator=gen)\n",
+    "    cnt = torch.bincount(ys, minlength=K).double() / N        # empirical pi_t\n",
+    "    # grad of (1/N) sum_i V log pi_theta(y_i) wrt logits at theta_t: V * (onehot - pi)\n",
+    "    gmc[x] = ((cnt * V[x])[:, None] * (torch.eye(K) - Pt[x][None, :])).sum(0) / m\n",
+    "err = (gmc - gJ).abs().max()\n",
+    "check('7. Monte Carlo RFT gradient (N=4e5 per prompt) vs grad J  (tol 3e-4)', err, 3e-4)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ba4f6acb",
+   "metadata": {},
+   "source": [
+    "## Verification table"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "3cb4ed0a",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:48:00.527272Z",
+     "iopub.status.busy": "2026-10-14T09:48:00.527101Z",
+     "iopub.status.idle": "2026-10-14T09:48:00.530845Z",
+     "shell.execute_reply": "2026-10-14T09:48:00.530126Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "identity                                                                                max abs err\n",
+      "1. grad of RFT objective at theta_t vs grad J(theta_t)                                    3.469e-18\n",
+      "2. ELBO identity  -KL(q||pi) = log p - KL(q||q*)                                          0.000e+00\n",
+      "2. tightness: -KL(q*||pi) = log p                                                         2.220e-16\n",
+      "3. p_{t+1}(x) = 1 on covered prompts (max |p-1|)                                          0.000e+00\n",
+      "3. p_{t+1}(x) = 0 on uncovered prompt (coverage failure)                                  0.000e+00\n",
+      "4. EM monotone: max decrease of F over 12 iterations (should be <= 0)                     0.000e+00\n",
+      "4. M-step solved to optimality (max |grad| at the last M-step)                            1.218e-09\n",
+      "4. ELBO(theta_t,q_t) = F(theta_t) at the final iterate                                    0.000e+00\n",
+      "4. equal-weight M-step objective: gradient at theta_t = grad F(theta_t) (weight 1/p_t per problem on grad p)    6.245e-17\n",
+      "4. pooled objective (ReST-EM Eq.3 / RFT): gradient at theta_t = grad mean_S p_t (weight 1 on grad p)    1.388e-17\n",
+      "6. non-monotone prompts exist under unnormalised weights (count of decreasing p, informational)    2.000e+01\n",
+      "5. E[filtered mean of f] = (1-(1-p)^n) E_q[f],  n=1                                       0.000e+00\n",
+      "5. E[filtered mean of f] = (1-(1-p)^n) E_q[f],  n=2                                       3.331e-16\n",
+      "5. E[filtered mean of f] = (1-(1-p)^n) E_q[f],  n=3                                       8.882e-16\n",
+      "5. E[filtered mean of f] = (1-(1-p)^n) E_q[f],  n=5                                       2.220e-15\n",
+      "5. pooled filtered average: E[(1/n) sum V f] = p E_q[f]                                   3.331e-16\n",
+      "7. Monte Carlo RFT gradient (N=4e5 per prompt) vs grad J  (tol 3e-4)                      1.941e-04\n",
+      "CHECKS_JSON:[[\"1. grad of RFT objective at theta_t vs grad J(theta_t)\", 3.469446951953614e-18, 1e-16], [\"2. ELBO identity  -KL(q||pi) = log p - KL(q||q*)\", 0.0, 1e-14], [\"2. tightness: -KL(q*||pi) = log p\", 2.220446049250313e-16, 1e-14], [\"3. p_{t+1}(x) = 1 on covered prompts (max |p-1|)\", 0.0, 1e-14], [\"3. p_{t+1}(x) = 0 on uncovered prompt (coverage failure)\", 0.0, 0.0], [\"4. EM monotone: max decrease of F over 12 iterations (should be <= 0)\", 0.0, 1e-10], [\"4. M-step solved to optimality (max |grad| at the last M-step)\", 1.2179568404172003e-09, 1e-08], [\"4. ELBO(theta_t,q_t) = F(theta_t) at the final iterate\", 0.0, 1e-12], [\"4. equal-weight M-step objective: gradient at theta_t = grad F(theta_t) (weight 1/p_t per problem on grad p)\", 6.245004513516506e-17, 1e-14], [\"4. pooled objective (ReST-EM Eq.3 / RFT): gradient at theta_t = grad mean_S p_t (weight 1 on grad p)\", 1.3877787807814457e-17, 1e-14], [\"6. non-monotone prompts exist under unnormalised weights (count of decreasing p, informational)\", 20.0, Infinity], [\"5. E[filtered mean of f] = (1-(1-p)^n) E_q[f],  n=1\", 0.0, 1e-14], [\"5. E[filtered mean of f] = (1-(1-p)^n) E_q[f],  n=2\", 3.3306690738754696e-16, 1e-14], [\"5. E[filtered mean of f] = (1-(1-p)^n) E_q[f],  n=3\", 8.881784197001252e-16, 1e-14], [\"5. E[filtered mean of f] = (1-(1-p)^n) E_q[f],  n=5\", 2.220446049250313e-15, 1e-14], [\"5. pooled filtered average: E[(1/n) sum V f] = p E_q[f]\", 3.3306690738754696e-16, 1e-14], [\"7. Monte Carlo RFT gradient (N=4e5 per prompt) vs grad J  (tol 3e-4)\", 0.00019412726116547693, 0.0003]]\n"
+     ]
+    }
+   ],
+   "source": [
+    "print()\n",
+    "print(f'{\"identity\":86s} {\"max abs err\":>12s}')\n",
+    "for n_, e_, t_ in CHECKS: print(f'{n_:86s} {e_:12.3e}')\n",
+    "print('CHECKS_JSON:' + json.dumps(CHECKS))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

code/notebooks/0203-pass-at-k.ipynb

@@ -0,0 +1,657 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "d6b3450b",
+   "metadata": {},
+   "source": [
+    "# 0203 -- pass@k: the unbiased estimator, its variance, and pass@1 vs pass@k\n",
+    "\n",
+    "Identities verified (exact rational arithmetic where possible, otherwise fp64):\n",
+    "1. `C(n,c) C(n-c,k) / C(n,k) = C(n-k,c)` and `E_c[ C(n-c,k)/C(n,k) ] = (1-p)^k` (exact `Fraction`).\n",
+    "2. Stable product form `prod_{i=n-c+1}^{n} (1-k/i)` equals the ratio of binomials; the naive float ratio overflows.\n",
+    "3. Variance by Hoeffding's U-statistic formula equals the exact binomial-sum variance; leading term `k^2 p (1-p)^(2k-1) / n`.\n",
+    "4. Plug-in `1-(1-p_hat)^k` is biased low: exact bias, and `-p(1-p)/n` for `k=2`.\n",
+    "5. `grad J_k = E_x[k (1-p)^(k-1) grad p]` against autograd.\n",
+    "6. A context-free strategy-mixing policy where the `J_1`-optimum and `J_32`-optimum differ (diversity vs sharpening).\n",
+    "7. The unbiased pass@k policy-gradient estimator: naive subset enumeration = closed-form coefficients; expectation = `grad J_k` exactly; leave-one-out form.\n",
+    "8. No unbiased estimator exists for `k > n` (Bernstein-basis projection residual)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "e6de1a23",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:57:46.427466Z",
+     "iopub.status.busy": "2026-10-14T08:57:46.427307Z",
+     "iopub.status.idle": "2026-10-14T08:57:48.231511Z",
+     "shell.execute_reply": "2026-10-14T08:57:48.230621Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n",
+    "import itertools, math, json\n",
+    "from fractions import Fraction as Fr\n",
+    "from math import comb\n",
+    "import numpy as np, torch\n",
+    "torch.set_default_dtype(torch.float64)\n",
+    "CHECKS = []\n",
+    "def check(name, err, tol):\n",
+    "    CHECKS.append((name, float(err), float(tol)))\n",
+    "    assert err <= tol, (name, float(err), tol)\n",
+    "    print(f'{name:80s} err={float(err):.3e}  (tol {tol:.0e})')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d785fb44",
+   "metadata": {},
+   "source": [
+    "## 1. Unbiasedness, exactly\n",
+    "For rational `p` the binomial expectation is a rational number, so the check is an *equality* of `Fraction`s."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "958d0931",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:57:48.233490Z",
+     "iopub.status.busy": "2026-10-14T08:57:48.233261Z",
+     "iopub.status.idle": "2026-10-14T08:57:48.338023Z",
+     "shell.execute_reply": "2026-10-14T08:57:48.337325Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1. E_c[1 - C(n-c,k)/C(n,k)] = 1-(1-p)^k, exact rational, all n<=40 tested        err=0.000e+00  (tol 0e+00)\n",
+      "1. identity C(n,c)C(n-c,k) = C(n,k)C(n-k,c) (violations over n<30)               err=0.000e+00  (tol 0e+00)\n"
+     ]
+    }
+   ],
+   "source": [
+    "def passk_unbiased_expectation(n, k, p):\n",
+    "    p = Fr(p)\n",
+    "    s = Fr(0)\n",
+    "    for c in range(n + 1):\n",
+    "        pc = comb(n, c) * p ** c * (1 - p) ** (n - c)\n",
+    "        s += pc * (1 - Fr(comb(n - c, k), comb(n, k)))          # comb(n-c,k)=0 when n-c<k\n",
+    "    return s\n",
+    "worst = Fr(0)\n",
+    "for n in (1, 2, 5, 12, 40):\n",
+    "    for k in range(1, n + 1):\n",
+    "        for p in (Fr(1, 7), Fr(1, 2), Fr(9, 10), Fr(1, 100)):\n",
+    "            worst = max(worst, abs(passk_unbiased_expectation(n, k, p) - (1 - (1 - p) ** k)))\n",
+    "check('1. E_c[1 - C(n-c,k)/C(n,k)] = 1-(1-p)^k, exact rational, all n<=40 tested', float(worst), 0.0)\n",
+    "bad = 0\n",
+    "for n in range(1, 30):\n",
+    "    for k in range(0, n + 1):\n",
+    "        for c in range(0, n + 1):\n",
+    "            if comb(n, c) * comb(n - c, k) * 1 != comb(n, k) * comb(n - k, c) * 1: bad += 1\n",
+    "check('1. identity C(n,c)C(n-c,k) = C(n,k)C(n-k,c) (violations over n<30)', bad, 0)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "374d4e73",
+   "metadata": {},
+   "source": [
+    "## 2. Stable implementation vs naive ratio of binomials\n",
+    "`C(n-c,k)/C(n,k) = prod_{i=n-c+1}^{n} (1 - k/i)`  (telescoping; derived in the note)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "8a3dc0ca",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:57:48.340113Z",
+     "iopub.status.busy": "2026-10-14T08:57:48.339952Z",
+     "iopub.status.idle": "2026-10-14T08:57:48.347641Z",
+     "shell.execute_reply": "2026-10-14T08:57:48.346855Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2. stable product form vs exact rational (max over grid incl. n=1000,k=999)      err=8.327e-17  (tol 1e-12)\n",
+      "naive float on n=6000,k=3000: OverflowError -> int too large to convert to float\n",
+      "stable value for the same input  : 1.0\n"
+     ]
+    }
+   ],
+   "source": [
+    "def pass_at_k_naive_float(n, c, k):\n",
+    "    return 1.0 - float(comb(n - c, k)) / float(comb(n, k))       # float(comb) overflows for big n,k\n",
+    "def pass_at_k_stable(n, c, k):\n",
+    "    if n - c < k: return 1.0\n",
+    "    return 1.0 - float(np.prod(1.0 - k / np.arange(n - c + 1, n + 1)))\n",
+    "def pass_at_k_exact(n, c, k): return 1 - Fr(comb(n - c, k), comb(n, k))\n",
+    "err = 0.0\n",
+    "for (n, k) in [(10, 3), (200, 10), (200, 100), (1000, 100), (1000, 999)]:\n",
+    "    for c in sorted(set([0, 1, 2, n // 7, n // 2, n - k, n - k + 1, n - 1, n])):\n",
+    "        if c > n: continue\n",
+    "        err = max(err, abs(pass_at_k_stable(n, c, k) - float(pass_at_k_exact(n, c, k))))\n",
+    "check('2. stable product form vs exact rational (max over grid incl. n=1000,k=999)', err, 1e-12)\n",
+    "try:\n",
+    "    v = pass_at_k_naive_float(6000, 100, 3000); print('naive float on n=6000,k=3000:', v)\n",
+    "    naive_failed = not np.isfinite(v)\n",
+    "except OverflowError as e:\n",
+    "    naive_failed = True; print('naive float on n=6000,k=3000: OverflowError ->', e)\n",
+    "assert naive_failed\n",
+    "print('stable value for the same input  :', pass_at_k_stable(6000, 100, 3000))\n",
+    "# the naive form silently loses precision / overflows; the stable form is a product of terms in (0,1)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d9320e2e",
+   "metadata": {},
+   "source": [
+    "## 3. Variance of the unbiased estimator\n",
+    "EXACT (naive): `Var = E[rho^2] - (1-p)^{2k}` summed over the binomial, `rho = C(n-c,k)/C(n,k)`.\n",
+    "HOEFFDING (efficient): the estimator is a U-statistic with kernel `h(S)=1{all k in S fail}`, and two random k-subsets overlapping in `j`\n",
+    "elements both fail with probability `(1-p)^{2k-j}`, so\n",
+    "`Var = sum_{j=1}^{k} w_j [ (1-p)^{2k-j} - (1-p)^{2k} ]`,  `w_j = C(k,j) C(n-k,k-j) / C(n,k)` (hypergeometric overlap pmf)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "3441dcee",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:57:48.349346Z",
+     "iopub.status.busy": "2026-10-14T08:57:48.349183Z",
+     "iopub.status.idle": "2026-10-14T08:57:48.377294Z",
+     "shell.execute_reply": "2026-10-14T08:57:48.376327Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "3. Hoeffding U-statistic variance = exact binomial-sum variance (exact Fractions) err=0.000e+00  (tol 1e-18)\n",
+      "n=200,k=8,p=0.3: exact Var=4.804470e-04  leading term k^2 p (1-p)^(2k-1)/n=4.557659e-04\n",
+      "3. exact variance vs leading term (relative error, n=200)                        err=5.415e-02  (tol 1e-01)\n"
+     ]
+    }
+   ],
+   "source": [
+    "def var_exact(n, k, p):\n",
+    "    p = Fr(p); m1 = Fr(0); m2 = Fr(0)\n",
+    "    for c in range(n + 1):\n",
+    "        pc = comb(n, c) * p ** c * (1 - p) ** (n - c); rho = Fr(comb(n - c, k), comb(n, k))\n",
+    "        m1 += pc * rho; m2 += pc * rho * rho\n",
+    "    return m2 - m1 * m1\n",
+    "def var_hoeffding(n, k, p):\n",
+    "    p = Fr(p); s = Fr(0)\n",
+    "    for j in range(1, k + 1):\n",
+    "        wj = Fr(comb(k, j) * comb(n - k, k - j), comb(n, k))\n",
+    "        s += wj * ((1 - p) ** (2 * k - j) - (1 - p) ** (2 * k))\n",
+    "    return s\n",
+    "worst = 0.0\n",
+    "for (n, k) in [(5, 2), (16, 4), (16, 8), (50, 10), (30, 30)]:\n",
+    "    for p in (Fr(1, 20), Fr(3, 10), Fr(1, 2), Fr(9, 10)):\n",
+    "        worst = max(worst, abs(float(var_exact(n, k, p) - var_hoeffding(n, k, p))))\n",
+    "check('3. Hoeffding U-statistic variance = exact binomial-sum variance (exact Fractions)', worst, 1e-18)\n",
+    "n, k, p = 200, 8, 0.3\n",
+    "lead = k ** 2 * p * (1 - p) ** (2 * k - 1) / n\n",
+    "print(f'n={n},k={k},p={p}: exact Var={float(var_exact(n,k,Fr(3,10))):.6e}  leading term k^2 p (1-p)^(2k-1)/n={lead:.6e}')\n",
+    "check('3. exact variance vs leading term (relative error, n=200)', abs(float(var_exact(n, k, Fr(3, 10))) / lead - 1), 0.12)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "dfad7109",
+   "metadata": {},
+   "source": [
+    "## 4. The plug-in estimator `1-(1-p_hat)^k` is biased low\n",
+    "By Jensen (`x -> (1-x)^k` is convex) `E[(1-p_hat)^k] >= (1-p)^k`. For `k=2`: `E(1-p_hat)^2 = (1-p)^2 + p(1-p)/n` exactly."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "6494c052",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:57:48.379114Z",
+     "iopub.status.busy": "2026-10-14T08:57:48.378897Z",
+     "iopub.status.idle": "2026-10-14T08:57:48.388166Z",
+     "shell.execute_reply": "2026-10-14T08:57:48.387120Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "4. k=2 exact bias of plug-in = -p(1-p)/n                                         err=0.000e+00  (tol 0e+00)\n",
+      "bias of plug-in pass@k (n=8,p=.3): {1: 0.0, 2: -0.02625, 4: -0.075460546875, 8: -0.10406236171669006}\n",
+      "4. k=1 plug-in is unbiased                                                       err=0.000e+00  (tol 0e+00)\n"
+     ]
+    }
+   ],
+   "source": [
+    "def plugin_expectation(n, k, p):\n",
+    "    p = Fr(p); return sum(comb(n, c) * p ** c * (1 - p) ** (n - c) * (1 - Fr(c, n)) ** k for c in range(n + 1))\n",
+    "n, p = 8, Fr(3, 10)\n",
+    "bias2 = (1 - plugin_expectation(n, 2, p)) - (1 - (1 - p) ** 2)       # E[plugin pass@2] - truth\n",
+    "check('4. k=2 exact bias of plug-in = -p(1-p)/n', float(abs(bias2 + p * (1 - p) / n)), 0.0)\n",
+    "print('bias of plug-in pass@k (n=8,p=.3):', {k_: float((1 - plugin_expectation(8, k_, p)) - (1 - (1 - p) ** k_)) for k_ in (1, 2, 4, 8)})\n",
+    "assert all((1 - plugin_expectation(8, k_, p)) - (1 - (1 - p) ** k_) <= 0 for k_ in range(1, 9))\n",
+    "check('4. k=1 plug-in is unbiased', float(abs((1 - plugin_expectation(8, 1, p)) - p)), 0.0)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1e4a9d5e",
+   "metadata": {},
+   "source": [
+    "## 5. Gradient of the pass@k objective: `k (1-p)^(k-1) grad p`\n",
+    "Tabular softmax per prompt, `J_k = mean_x [1-(1-p_x)^k]`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "f7df0b58",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:57:48.390043Z",
+     "iopub.status.busy": "2026-10-14T08:57:48.389832Z",
+     "iopub.status.idle": "2026-10-14T08:57:48.487130Z",
+     "shell.execute_reply": "2026-10-14T08:57:48.486143Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "5. grad J_k = E_x[k (1-p)^{k-1} grad p] vs autograd, k in {1,2,8,64}             err=6.939e-18  (tol 1e-16)\n",
+      "weights w_k(p)=k(1-p)^{k-1}:\n",
+      "  k=  1 [1. 1. 1. 1. 1. 1.]\n",
+      "  k=  4 [3.8812 3.4295 2.916  1.372  0.5    0.004 ]\n",
+      "  k= 16 [1.37609e+01 7.41270e+00 3.29430e+00 7.60000e-02 5.00000e-04 0.00000e+00]\n",
+      "  k= 64 [33.978   2.5279  0.0838  0.      0.      0.    ]\n"
+     ]
+    }
+   ],
+   "source": [
+    "torch.manual_seed(0)\n",
+    "Kout, Xn = 9, 7\n",
+    "Vtab = (torch.rand(Xn, Kout) < 0.25).double(); Vtab[0, :] = 0; Vtab[0, 3] = 1\n",
+    "z = (torch.randn(Xn, Kout) * 0.8).requires_grad_(True)\n",
+    "def p_x(z): return (torch.softmax(z, -1) * Vtab).sum(-1)\n",
+    "worst = 0.0\n",
+    "for k in (1, 2, 8, 64):\n",
+    "    z.grad = None\n",
+    "    Jk = (1 - (1 - p_x(z)) ** k).mean(); Jk.backward(); g_auto = z.grad.clone()\n",
+    "    z2 = z.detach().clone().requires_grad_(True)\n",
+    "    pz = p_x(z2); pz.sum().backward()                       # d p_x / dz row-wise (each p_x depends only on its row)\n",
+    "    gp = z2.grad                                            # row x holds grad p_x\n",
+    "    g_hand = (k * (1 - pz.detach()) ** (k - 1))[:, None] * gp / Xn\n",
+    "    worst = max(worst, float((g_auto - g_hand).abs().max()))\n",
+    "check('5. grad J_k = E_x[k (1-p)^{k-1} grad p] vs autograd, k in {1,2,8,64}', worst, 1e-16)\n",
+    "pv = np.array([0.01, 0.05, 0.1, 0.3, 0.5, 0.9])\n",
+    "print('weights w_k(p)=k(1-p)^{k-1}:')\n",
+    "for k in (1, 4, 16, 64): print(f'  k={k:3d}', np.round(k * (1 - pv) ** (k - 1), 4))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6a670727",
+   "metadata": {},
+   "source": [
+    "## 6. Why optimising pass@1 and pass@k pull in different directions\n",
+    "Within one prompt `pass@k = 1-(1-p)^k` is increasing in `p`, so the two objectives *agree* on that prompt. They disagree **across prompts**\n",
+    "through the weights `1` vs `k(1-p)^{k-1}`. A context-free \"strategy\" policy makes this concrete: `pi` over 4 strategies; strategy 0 is a\n",
+    "broad generalist that solves 8 \"easy\" prompts with prob 0.6 and nothing else; strategies 1..3 each solve one of 4 hard prompts... with prob 0.3\n",
+    "(strategy 1 -> hard prompts 8, 9; strategy 2 -> 10; strategy 3 -> 11) and 0.1 on easy ones."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "0e976d80",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:57:48.488964Z",
+     "iopub.status.busy": "2026-10-14T08:57:48.488802Z",
+     "iopub.status.idle": "2026-10-14T08:57:53.482658Z",
+     "shell.execute_reply": "2026-10-14T08:57:53.481750Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "argmax J_1 : pi = [1. 0. 0. 0.]  entropy 0.0\n",
+      "argmax J_32: pi = [0.111  0.3414 0.2738 0.2738]  entropy 1.3202\n",
+      "J_1 at th1=0.4000  J_1 at th32=0.1344\n",
+      "J_32 at th1=0.6667  J_32 at th32=0.9810\n",
+      "k=1: |grad| at optimum = 2.12e-07\n",
+      "k=32: |grad| at optimum = 4.82e-05\n",
+      "6. stationarity of both optima (max |grad J_k| at the returned theta)            err=4.823e-05  (tol 1e-03)\n",
+      "6. cross-inferiority gaps: min(J_1(th1)-J_1(th32), J_32(th32)-J_32(th1)) > 0 (reported as 0 if true) err=0.000e+00  (tol 0e+00)\n"
+     ]
+    }
+   ],
+   "source": [
+    "Xn = 12\n",
+    "S = torch.zeros(Xn, 4)\n",
+    "S[:8, 0] = 0.6; S[:8, 1:] = 0.1\n",
+    "S[8, 1] = 0.3; S[9, 1] = 0.3; S[10, 2] = 0.3; S[11, 3] = 0.3\n",
+    "def J_k(theta, k):\n",
+    "    pi = torch.softmax(theta, 0); p = S @ pi\n",
+    "    return (1 - (1 - p) ** k).mean()\n",
+    "def opt(k, steps=6000, lr=0.3, seed=0):\n",
+    "    th = torch.zeros(4, requires_grad=True)\n",
+    "    o = torch.optim.Adam([th], lr=lr)\n",
+    "    for _ in range(steps):\n",
+    "        o.zero_grad(); (-J_k(th, k)).backward(); o.step()\n",
+    "    return th.detach()\n",
+    "th1, th32 = opt(1), opt(32)\n",
+    "H = lambda th: float(-(torch.softmax(th, 0) * torch.log_softmax(th, 0)).sum())\n",
+    "print('argmax J_1 : pi =', np.round(torch.softmax(th1, 0).numpy(), 4), ' entropy', round(H(th1), 4))\n",
+    "print('argmax J_32: pi =', np.round(torch.softmax(th32, 0).numpy(), 4), ' entropy', round(H(th32), 4))\n",
+    "print(f'J_1 at th1={float(J_k(th1,1)):.4f}  J_1 at th32={float(J_k(th32,1)):.4f}')\n",
+    "print(f'J_32 at th1={float(J_k(th1,32)):.4f}  J_32 at th32={float(J_k(th32,32)):.4f}')\n",
+    "assert J_k(th1, 1) > J_k(th32, 1) + 1e-3 and J_k(th32, 32) > J_k(th1, 32) + 1e-3 and H(th32) > H(th1)\n",
+    "# Check that the optimiser reached a stationary point (max |grad| of J_k at the returned theta)\n",
+    "gmax = []\n",
+    "for k, th in ((1, th1), (32, th32)):\n",
+    "    t = th.clone().requires_grad_(True); J_k(t, k).backward()\n",
+    "    gmax.append(float(t.grad.abs().max())); print(f'k={k}: |grad| at optimum = {gmax[-1]:.2e}')\n",
+    "check('6. stationarity of both optima (max |grad J_k| at the returned theta)', max(gmax), 1e-3)\n",
+    "check('6. cross-inferiority gaps: min(J_1(th1)-J_1(th32), J_32(th32)-J_32(th1)) > 0 (reported as 0 if true)', 0.0 if min(float(J_k(th1,1)-J_k(th32,1)), float(J_k(th32,32)-J_k(th1,32))) > 0 else 1.0, 0.0)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "75d59643",
+   "metadata": {},
+   "source": [
+    "## 7. The unbiased pass@k *gradient* estimator (PKPO-style), exact enumeration\n",
+    "Tabular softmax over `K=4` outcomes with binary reward vector `r`, `n=5` samples, `k=3`. For a tuple `(y_1..y_n)`:\n",
+    "* NAIVE: `g = (1/C(n,k)) sum_S max_{i in S} r_i * sum_{i in S} grad log pi(y_i)` (all `C(n,k)` subsets).\n",
+    "* CLOSED FORM: `g = sum_i w_i grad log pi(y_i)` with `w_i = (k/n)[ rho_{-i} + r_i (1 - rho_{-i}) ]`, where `rho_{-i}=1-C(n-1-c_{-i},k-1)/C(n-1,k-1)`\n",
+    "  is the unbiased pass@(k-1) estimate from the other `n-1` samples, `c_{-i}` the number correct among them.\n",
+    "* LEAVE-ONE-OUT form: subtract the baseline `(k/n) rho_{-i}` (independent of `y_i`): coefficient `(k/n) r_i (1-rho_{-i})`.\n",
+    "`grad J_k = k (1-p)^{k-1} grad p` is the target."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "556c657b",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:57:53.484341Z",
+     "iopub.status.busy": "2026-10-14T08:57:53.484099Z",
+     "iopub.status.idle": "2026-10-14T08:57:54.020052Z",
+     "shell.execute_reply": "2026-10-14T08:57:54.019291Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "7. closed-form coefficients = naive subset average, every tuple (max abs diff)   err=4.441e-16  (tol 1e-14)\n",
+      "7. E[naive estimator] = grad of 1-(1-p)^k  (exact over 4^5 tuples)               err=4.302e-16  (tol 1e-15)\n",
+      "7. E[closed-form] = grad J_k                                                     err=4.302e-16  (tol 1e-15)\n",
+      "7. E[leave-one-out form] = grad J_k  (baseline (k/n) rho_{-i} is independent of y_i) err=6.661e-16  (tol 1e-15)\n",
+      "total variance: closed-form (no baseline)=0.44125   leave-one-out=0.13981   ratio=0.317\n",
+      "7. k=1 reduces to REINFORCE sum_i (r_i/n) grad log pi(y_i)                       err=1.110e-16  (tol 1e-15)\n"
+     ]
+    }
+   ],
+   "source": [
+    "Kc, n, k = 4, 5, 3\n",
+    "r = torch.tensor([0., 1., 0., 1.])\n",
+    "theta = torch.tensor([0.3, -0.5, 0.9, -1.1], requires_grad=True)\n",
+    "pi = torch.softmax(theta, 0)\n",
+    "pk_obj = 1 - (1 - (pi * r).sum()) ** k\n",
+    "pk_obj.backward(); g_target = theta.grad.clone()\n",
+    "pi_d = pi.detach()\n",
+    "score = torch.eye(Kc) - pi_d[None, :]               # row y: grad_theta log pi(y)\n",
+    "def rho_minus(c_other, n_, k_):                     # unbiased pass@(k-1) from n-1 samples with c_other correct (k_=k)\n",
+    "    kk = k_ - 1\n",
+    "    if kk == 0: return 0.0\n",
+    "    return 1.0 - comb(n_ - 1 - c_other, kk) / comb(n_ - 1, kk)\n",
+    "E_naive = torch.zeros(Kc); E_closed = torch.zeros(Kc); E_loo = torch.zeros(Kc)\n",
+    "maxdiff = 0.0; var_closed = torch.zeros(Kc); var_loo = torch.zeros(Kc)\n",
+    "tuples = list(itertools.product(range(Kc), repeat=n))\n",
+    "second_closed = torch.zeros(Kc); second_loo = torch.zeros(Kc)\n",
+    "for tup in tuples:\n",
+    "    prob = float(np.prod([float(pi_d[y]) for y in tup]))\n",
+    "    rr = [float(r[y]) for y in tup]\n",
+    "    # naive: all subsets\n",
+    "    g1 = torch.zeros(Kc)\n",
+    "    for Ssub in itertools.combinations(range(n), k):\n",
+    "        mx = max(rr[i] for i in Ssub)\n",
+    "        g1 += mx * sum(score[tup[i]] for i in Ssub)\n",
+    "    g1 /= comb(n, k)\n",
+    "    # closed form\n",
+    "    g2 = torch.zeros(Kc); g3 = torch.zeros(Kc); c_tot = sum(rr)\n",
+    "    for i in range(n):\n",
+    "        c_other = c_tot - rr[i]\n",
+    "        rho = rho_minus(int(c_other), n, k)\n",
+    "        w = (k / n) * (rho + rr[i] * (1 - rho))\n",
+    "        g2 += w * score[tup[i]]\n",
+    "        g3 += (k / n) * rr[i] * (1 - rho) * score[tup[i]]\n",
+    "    maxdiff = max(maxdiff, float((g1 - g2).abs().max()))\n",
+    "    E_naive += prob * g1; E_closed += prob * g2; E_loo += prob * g3\n",
+    "    second_closed += prob * g2 ** 2; second_loo += prob * g3 ** 2\n",
+    "check('7. closed-form coefficients = naive subset average, every tuple (max abs diff)', maxdiff, 1e-14)\n",
+    "check('7. E[naive estimator] = grad of 1-(1-p)^k  (exact over 4^5 tuples)', (E_naive - g_target).abs().max(), 1e-15)\n",
+    "check('7. E[closed-form] = grad J_k', (E_closed - g_target).abs().max(), 1e-15)\n",
+    "check('7. E[leave-one-out form] = grad J_k  (baseline (k/n) rho_{-i} is independent of y_i)', (E_loo - g_target).abs().max(), 1e-15)\n",
+    "v_closed = float((second_closed - E_closed ** 2).sum()); v_loo = float((second_loo - E_loo ** 2).sum())\n",
+    "print(f'total variance: closed-form (no baseline)={v_closed:.5f}   leave-one-out={v_loo:.5f}   ratio={v_loo/v_closed:.3f}')\n",
+    "assert v_loo < v_closed\n",
+    "# k=1 reduces to REINFORCE without baseline: w_i = r_i/n\n",
+    "kk = 1; maxd = 0.0\n",
+    "for tup in tuples[:200]:\n",
+    "    rr = [float(r[y]) for y in tup]\n",
+    "    g1 = sum(max(rr[i] for i in (j,)) * score[tup[j]] for j in range(n)) / n\n",
+    "    g2 = sum((kk / n) * (0.0 + rr[i] * (1 - 0.0)) * score[tup[i]] for i in range(n))\n",
+    "    maxd = max(maxd, float((g1 - g2).abs().max()))\n",
+    "check('7. k=1 reduces to REINFORCE sum_i (r_i/n) grad log pi(y_i)', maxd, 1e-15)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "65ffaf66",
+   "metadata": {},
+   "source": [
+    "## 8. No unbiased estimator for `k > n` (Kolmogorov): Bernstein-basis residual\n",
+    "Unbiased estimators of `f(p)` from `n` Bernoulli draws are exactly `f(p) = sum_c a_c C(n,c) p^c (1-p)^{n-c}`, i.e. polynomials of degree `<= n`.\n",
+    "Project `f(p) = 1-(1-p)^k` on the degree-`n` Bernstein basis: residual 0 iff `k <= n`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "476de431",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:57:54.021866Z",
+     "iopub.status.busy": "2026-10-14T08:57:54.021703Z",
+     "iopub.status.idle": "2026-10-14T08:57:54.031118Z",
+     "shell.execute_reply": "2026-10-14T08:57:54.030302Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{(4, 3): 1.2212453270876722e-15, (4, 4): 9.558326352632207e-16, (4, 5): 0.003948447395839216, (4, 8): 0.043274555276372625, (10, 10): 2.617697725249002e-15, (10, 11): 1.3790524725575724e-06}\n",
+      "8. residual for k<=n (max over tested) ~ 0                                       err=2.618e-15  (tol 1e-09)\n"
+     ]
+    }
+   ],
+   "source": [
+    "pg = np.linspace(0, 1, 2001)\n",
+    "def bern_resid(n_, k_):\n",
+    "    B = np.stack([comb(n_, c) * pg ** c * (1 - pg) ** (n_ - c) for c in range(n_ + 1)], 1)\n",
+    "    f = 1 - (1 - pg) ** k_\n",
+    "    a, *_ = np.linalg.lstsq(B, f, rcond=None)\n",
+    "    return float(np.abs(B @ a - f).max())\n",
+    "res = {(n_, k_): bern_resid(n_, k_) for n_, k_ in [(4, 3), (4, 4), (4, 5), (4, 8), (10, 10), (10, 11)]}\n",
+    "print(res)\n",
+    "check('8. residual for k<=n (max over tested) ~ 0', max(res[(4, 3)], res[(4, 4)], res[(10, 10)]), 1e-9)\n",
+    "assert res[(4, 5)] > 1e-3 and res[(10, 11)] > 1e-6\n",
+    "CHECKS.append(('8. residual for k=n+1 is strictly positive (min over tested; informational)', float(min(res[(4,5)], res[(10,11)])), float('inf')))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "19e22270",
+   "metadata": {},
+   "source": [
+    "## 9. Benchmark-level sampling variance: within-problem part\n",
+    "`N=244` problems with `p_x ~ Beta(0.5, 1)`, `n=16` samples per problem, `k=8`. For *fixed* problems the variance of the benchmark mean is\n",
+    "`(1/N^2) sum_x Var_x`, with `Var_x` from section 3. Monte Carlo agrees."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "4691ab9e",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:57:54.032552Z",
+     "iopub.status.busy": "2026-10-14T08:57:54.032377Z",
+     "iopub.status.idle": "2026-10-14T08:58:10.317977Z",
+     "shell.execute_reply": "2026-10-14T08:58:10.317179Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "truth=0.67114  MC mean=0.67109 (+-0.00006)  MC var=1.3361e-04  theory var=1.3230e-04\n",
+      "9. estimator is unbiased at benchmark level (|MC mean - truth| in s.e.)          err=9.278e-01  (tol 4e+00)\n",
+      "9. within-problem variance: MC vs theory (relative)                              err=9.909e-03  (tol 3e-02)\n"
+     ]
+    }
+   ],
+   "source": [
+    "rg = np.random.RandomState(5)\n",
+    "N, n, k = 244, 16, 8\n",
+    "px = rg.beta(0.5, 1.0, size=N)\n",
+    "truth = float((1 - (1 - px) ** k).mean())\n",
+    "var_theory = sum(float(var_exact(n, k, Fr(float(pv)).limit_denominator(10 ** 9))) for pv in px) / N ** 2\n",
+    "reps = 40000\n",
+    "c = rg.binomial(n, px[None, :], size=(reps, N))\n",
+    "def est_from_c(c):\n",
+    "    out = np.ones_like(c, dtype=float)\n",
+    "    for cc in range(0, n + 1):\n",
+    "        out[c == cc] = pass_at_k_stable(n, cc, k)\n",
+    "    return out.mean(1)\n",
+    "bench = est_from_c(c)\n",
+    "se_mean = bench.std() / math.sqrt(reps)\n",
+    "print(f'truth={truth:.5f}  MC mean={bench.mean():.5f} (+-{se_mean:.5f})  MC var={bench.var():.4e}  theory var={var_theory:.4e}')\n",
+    "check('9. estimator is unbiased at benchmark level (|MC mean - truth| in s.e.)', abs(bench.mean() - truth) / se_mean, 4.0)\n",
+    "check('9. within-problem variance: MC vs theory (relative)', abs(bench.var() / var_theory - 1), 0.03)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "83602bfb",
+   "metadata": {},
+   "source": [
+    "## Verification table"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "86c128e2",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T08:58:10.319495Z",
+     "iopub.status.busy": "2026-10-14T08:58:10.319334Z",
+     "iopub.status.idle": "2026-10-14T08:58:10.322873Z",
+     "shell.execute_reply": "2026-10-14T08:58:10.322203Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "identity                                                                                  max abs err\n",
+      "1. E_c[1 - C(n-c,k)/C(n,k)] = 1-(1-p)^k, exact rational, all n<=40 tested                   0.000e+00\n",
+      "1. identity C(n,c)C(n-c,k) = C(n,k)C(n-k,c) (violations over n<30)                          0.000e+00\n",
+      "2. stable product form vs exact rational (max over grid incl. n=1000,k=999)                 8.327e-17\n",
+      "3. Hoeffding U-statistic variance = exact binomial-sum variance (exact Fractions)           0.000e+00\n",
+      "3. exact variance vs leading term (relative error, n=200)                                   5.415e-02\n",
+      "4. k=2 exact bias of plug-in = -p(1-p)/n                                                    0.000e+00\n",
+      "4. k=1 plug-in is unbiased                                                                  0.000e+00\n",
+      "5. grad J_k = E_x[k (1-p)^{k-1} grad p] vs autograd, k in {1,2,8,64}                        6.939e-18\n",
+      "6. stationarity of both optima (max |grad J_k| at the returned theta)                       4.823e-05\n",
+      "6. cross-inferiority gaps: min(J_1(th1)-J_1(th32), J_32(th32)-J_32(th1)) > 0 (reported as 0 if true)    0.000e+00\n",
+      "7. closed-form coefficients = naive subset average, every tuple (max abs diff)              4.441e-16\n",
+      "7. E[naive estimator] = grad of 1-(1-p)^k  (exact over 4^5 tuples)                          4.302e-16\n",
+      "7. E[closed-form] = grad J_k                                                                4.302e-16\n",
+      "7. E[leave-one-out form] = grad J_k  (baseline (k/n) rho_{-i} is independent of y_i)        6.661e-16\n",
+      "7. k=1 reduces to REINFORCE sum_i (r_i/n) grad log pi(y_i)                                  1.110e-16\n",
+      "8. residual for k<=n (max over tested) ~ 0                                                  2.618e-15\n",
+      "8. residual for k=n+1 is strictly positive (min over tested; informational)                 1.379e-06\n",
+      "9. estimator is unbiased at benchmark level (|MC mean - truth| in s.e.)                     9.278e-01\n",
+      "9. within-problem variance: MC vs theory (relative)                                         9.909e-03\n",
+      "CHECKS_JSON:[[\"1. E_c[1 - C(n-c,k)/C(n,k)] = 1-(1-p)^k, exact rational, all n<=40 tested\", 0.0, 0.0], [\"1. identity C(n,c)C(n-c,k) = C(n,k)C(n-k,c) (violations over n<30)\", 0.0, 0.0], [\"2. stable product form vs exact rational (max over grid incl. n=1000,k=999)\", 8.326672684688674e-17, 1e-12], [\"3. Hoeffding U-statistic variance = exact binomial-sum variance (exact Fractions)\", 0.0, 1e-18], [\"3. exact variance vs leading term (relative error, n=200)\", 0.0541530729940769, 0.12], [\"4. k=2 exact bias of plug-in = -p(1-p)/n\", 0.0, 0.0], [\"4. k=1 plug-in is unbiased\", 0.0, 0.0], [\"5. grad J_k = E_x[k (1-p)^{k-1} grad p] vs autograd, k in {1,2,8,64}\", 6.938893903907228e-18, 1e-16], [\"6. stationarity of both optima (max |grad J_k| at the returned theta)\", 4.8228746018059434e-05, 0.001], [\"6. cross-inferiority gaps: min(J_1(th1)-J_1(th32), J_32(th32)-J_32(th1)) > 0 (reported as 0 if true)\", 0.0, 0.0], [\"7. closed-form coefficients = naive subset average, every tuple (max abs diff)\", 4.440892098500626e-16, 1e-14], [\"7. E[naive estimator] = grad of 1-(1-p)^k  (exact over 4^5 tuples)\", 4.3021142204224816e-16, 1e-15], [\"7. E[closed-form] = grad J_k\", 4.3021142204224816e-16, 1e-15], [\"7. E[leave-one-out form] = grad J_k  (baseline (k/n) rho_{-i} is independent of y_i)\", 6.661338147750939e-16, 1e-15], [\"7. k=1 reduces to REINFORCE sum_i (r_i/n) grad log pi(y_i)\", 1.1102230246251565e-16, 1e-15], [\"8. residual for k<=n (max over tested) ~ 0\", 2.617697725249002e-15, 1e-09], [\"8. residual for k=n+1 is strictly positive (min over tested; informational)\", 1.3790524725575724e-06, Infinity], [\"9. estimator is unbiased at benchmark level (|MC mean - truth| in s.e.)\", 0.9277896993916872, 4.0], [\"9. within-problem variance: MC vs theory (relative)\", 0.009909418241782264, 0.03]]\n"
+     ]
+    }
+   ],
+   "source": [
+    "print()\n",
+    "print(f'{\"identity\":88s} {\"max abs err\":>12s}')\n",
+    "for n_, e_, t_ in CHECKS: print(f'{n_:88s} {e_:12.3e}')\n",
+    "print('CHECKS_JSON:' + json.dumps(CHECKS))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

code/notebooks/0204-self-correction-multiturn-mdp.ipynb

@@ -0,0 +1,505 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "bc47f601",
+   "metadata": {},
+   "source": [
+    "# 0204 -- Self-correction as a multi-turn MDP: masks, per-turn credit, baselines\n",
+    "\n",
+    "A two-turn tabular environment small enough to *enumerate every episode and every group of sampled episodes exactly*.\n",
+    "Turn 1: `y1 ~ pi1` (K1=3). The verifier returns `r1(y1)` and a feedback token `f = F(y1)` (deterministic).\n",
+    "Turn 2: `y2 ~ pi2(.|y1)` (K2=3), context = `(x, y1, f)`. Verifier returns `r2(y1,y2)`. Episode return `R = r1 + r2`.\n",
+    "Parameters: `theta1 in R^3`, `theta2 in R^{3x3}` (one softmax row per turn-2 context).\n",
+    "\n",
+    "Identities verified:\n",
+    "1. Trajectory factorisation: the environment factor drops out; masked REINFORCE is unbiased, `E = grad J`.\n",
+    "2. Past rewards are independent of later actions: `E[r1 * grad log pi2(y2|c2)] = 0` (reward-to-go is unbiased, with lower variance).\n",
+    "3. Including feedback tokens in the loss adds a non-zero bias `E[R grad log pi(f|.)]` when the feedback shares the output softmax.\n",
+    "4. GRPO's mean-including-self baseline is exactly `(G-1)/G` times the leave-one-out one, so its gradient is `(G-1)/G * grad J`.\n",
+    "5. Leave-one-out (episode-level) is unbiased; a baseline that depends on the sample's own outcome is biased.\n",
+    "6. Two-level (tree) leave-one-out with per-turn baselines is unbiased; variance compared at equal episode budget.\n",
+    "7. Training only on a fixed (stale) distribution of turn-2 contexts is biased by `sum (mu - pi1) grad Q`, and drops the turn-1 term."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "343280c0",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:01:50.331677Z",
+     "iopub.status.busy": "2026-10-14T09:01:50.331456Z",
+     "iopub.status.idle": "2026-10-14T09:01:52.235734Z",
+     "shell.execute_reply": "2026-10-14T09:01:52.234848Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "J = 1.0181152430262448\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/vignesh/.opencode/tmp/ipykernel_471603/844665086.py:31: UserWarning: Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.\n",
+      "Consider using tensor.detach() first. (Triggered internally at /pytorch/torch/csrc/autograd/generated/python_variable_methods.cpp:838.)\n",
+      "  print('J =', float(J))\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n",
+    "import itertools, math, json\n",
+    "import numpy as np, torch\n",
+    "torch.set_default_dtype(torch.float64)\n",
+    "CHECKS = []\n",
+    "def check(name, err, tol):\n",
+    "    CHECKS.append((name, float(err), float(tol)))\n",
+    "    assert err <= tol, (name, float(err), tol)\n",
+    "    print(f'{name:86s} err={float(err):.3e}  (tol {tol:.0e})')\n",
+    "\n",
+    "K1, K2 = 3, 3\n",
+    "g = np.random.RandomState(4)\n",
+    "r1 = np.array([0., 0., 1.])                      # turn-1 success table\n",
+    "r2 = np.array([[1, 0, 0], [0, 1, 0], [0, 1, 1]], float)   # turn-2 success depends on (y1, y2)\n",
+    "Fmap = np.array([0, 1, 2])                        # feedback token id as a function of y1 (shares the K2-sized vocabulary)\n",
+    "theta1_0 = torch.tensor([0.2, -0.4, 0.5])\n",
+    "theta2_0 = torch.tensor(g.randn(K1, K2) * 0.7)\n",
+    "nparam = K1 + K1 * K2\n",
+    "\n",
+    "def pis(t1, t2): return torch.softmax(t1, 0), torch.softmax(t2, -1)\n",
+    "def leaves(t1, t2):\n",
+    "    p1, p2 = pis(t1, t2)\n",
+    "    return {(a, b): float(p1[a] * p2[a, b]) for a in range(K1) for b in range(K2)}\n",
+    "def J_fn(t1, t2):\n",
+    "    p1, p2 = pis(t1, t2)\n",
+    "    return sum(p1[a] * p2[a, b] * (r1[a] + r2[a, b]) for a in range(K1) for b in range(K2))\n",
+    "t1 = theta1_0.clone().requires_grad_(True); t2 = theta2_0.clone().requires_grad_(True)\n",
+    "J = J_fn(t1, t2); J.backward()\n",
+    "gradJ = torch.cat([t1.grad, t2.grad.flatten()]).numpy()\n",
+    "print('J =', float(J))\n",
+    "p1n, p2n = [x.detach().numpy() for x in pis(theta1_0, theta2_0)]\n",
+    "P = {(a, b): p1n[a] * p2n[a, b] for a in range(K1) for b in range(K2)}\n",
+    "def s1(a):        # score of turn-1 token: e_a - pi1 in theta1 block\n",
+    "    v = np.zeros(nparam); e = np.zeros(K1); e[a] = 1; v[:K1] = e - p1n; return v\n",
+    "def s2(a, b):     # score of turn-2 token: e_b - pi2[a] in the a-th row of theta2\n",
+    "    v = np.zeros(nparam); e = np.zeros(K2); e[b] = 1; v[K1 + a * K2: K1 + (a + 1) * K2] = e - p2n[a]; return v\n",
+    "def sF(a):        # score of the *feedback token* f=F(a) if it were scored by the same softmax head as turn 2, context y1=a\n",
+    "    v = np.zeros(nparam); e = np.zeros(K2); e[Fmap[a]] = 1; v[K1 + a * K2: K1 + (a + 1) * K2] = e - p2n[a]; return v\n",
+    "Rret = {(a, b): r1[a] + r2[a, b] for a in range(K1) for b in range(K2)}\n",
+    "L = list(P.keys())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d24cef3e",
+   "metadata": {},
+   "source": [
+    "## 1. Masked REINFORCE is unbiased\n",
+    "Trajectory probability: `P(tau) = pi1(y1) * E(f|y1) * pi2(y2|y1,f)`. The environment factor `E` does not depend on theta, so\n",
+    "`grad log P(tau) = grad log pi1(y1) + grad log pi2(y2|c2)` -- only policy-generated tokens appear."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "903afa0a",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:01:52.237467Z",
+     "iopub.status.busy": "2026-10-14T09:01:52.237201Z",
+     "iopub.status.idle": "2026-10-14T09:01:52.241076Z",
+     "shell.execute_reply": "2026-10-14T09:01:52.240235Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1. E[R (score_turn1 + score_turn2)] = grad J (exact enumeration of 9 episodes)         err=8.327e-17  (tol 1e-15)\n"
+     ]
+    }
+   ],
+   "source": [
+    "E_mask = sum(P[l] * Rret[l] * (s1(l[0]) + s2(*l)) for l in L)\n",
+    "check('1. E[R (score_turn1 + score_turn2)] = grad J (exact enumeration of 9 episodes)', np.abs(E_mask - gradJ).max(), 1e-15)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "851be585",
+   "metadata": {},
+   "source": [
+    "## 2. Causality: a past reward carries no information about a later action\n",
+    "`E[ r1(y1) * grad log pi2(y2|y1) ] = 0` because, given `y1`, `E_{y2}[grad log pi2] = 0`. So `r1` may be dropped from the turn-2 coefficient."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "8c6af3a3",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:01:52.242434Z",
+     "iopub.status.busy": "2026-10-14T09:01:52.242279Z",
+     "iopub.status.idle": "2026-10-14T09:01:52.248993Z",
+     "shell.execute_reply": "2026-10-14T09:01:52.248126Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2. E[r1 * score_turn2] = 0                                                             err=1.388e-17  (tol 1e-16)\n",
+      "2. reward-to-go estimator: E = grad J                                                  err=6.939e-17  (tol 1e-15)\n",
+      "single-episode total variance: full return=1.6566  reward-to-go=0.9474\n"
+     ]
+    }
+   ],
+   "source": [
+    "E_past = sum(P[l] * r1[l[0]] * s2(*l) for l in L)\n",
+    "check('2. E[r1 * score_turn2] = 0', np.abs(E_past).max(), 1e-16)\n",
+    "E_rtg = sum(P[l] * (Rret[l] * s1(l[0]) + r2[l] * s2(*l)) for l in L)\n",
+    "check('2. reward-to-go estimator: E = grad J', np.abs(E_rtg - gradJ).max(), 1e-15)\n",
+    "var_full = sum(P[l] * ((Rret[l] * (s1(l[0]) + s2(*l))) ** 2).sum() for l in L) - (gradJ ** 2).sum()\n",
+    "var_rtg = sum(P[l] * ((Rret[l] * s1(l[0]) + r2[l] * s2(*l)) ** 2).sum() for l in L) - (gradJ ** 2).sum()\n",
+    "print(f'single-episode total variance: full return={var_full:.4f}  reward-to-go={var_rtg:.4f}')\n",
+    "assert var_rtg < var_full"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ba273ae5",
+   "metadata": {},
+   "source": [
+    "## 3. Putting the feedback token in the loss biases the gradient\n",
+    "If the loss also includes `log pi(f|context)` for the environment-written feedback token (same output softmax), the expected extra term is\n",
+    "`E[R grad log pi2(f|y1)]`, generally non-zero."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "5f223ef3",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:01:52.250555Z",
+     "iopub.status.busy": "2026-10-14T09:01:52.250399Z",
+     "iopub.status.idle": "2026-10-14T09:01:52.255258Z",
+     "shell.execute_reply": "2026-10-14T09:01:52.254500Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "||bias from unmasked feedback tokens|| = 0.679933398028168  vs ||grad J|| = 0.47960612146460163\n",
+      "with constant R the same term is 0.5000121536758657 (non-zero: f is not distributed as pi2, so even the baseline identity fails)\n"
+     ]
+    }
+   ],
+   "source": [
+    "E_fb = sum(P[l] * Rret[l] * sF(l[0]) for l in L)\n",
+    "print('||bias from unmasked feedback tokens|| =', float(np.linalg.norm(E_fb)), ' vs ||grad J|| =', float(np.linalg.norm(gradJ)))\n",
+    "assert np.linalg.norm(E_fb) > 1e-3\n",
+    "CHECKS.append(('3. norm of the bias from including feedback tokens in the loss (informational, must be > 0)', float(np.linalg.norm(E_fb)), float('inf')))\n",
+    "# without a reward correlation (R constant) the term vanishes: E[score of a fixed-distribution token] = 0 only for tokens sampled from pi.\n",
+    "E_fb_const = sum(P[l] * 1.0 * sF(l[0]) for l in L)\n",
+    "print('with constant R the same term is', float(np.linalg.norm(E_fb_const)), '(non-zero: f is not distributed as pi2, so even the baseline identity fails)')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0b2c937c",
+   "metadata": {},
+   "source": [
+    "## 4--6. Group estimators, exactly\n",
+    "`G` i.i.d. episodes. Enumerate all `9^G` groups with their probabilities. Estimators:\n",
+    "* `E0`  no baseline, full return; `E0rtg` reward-to-go\n",
+    "* `E1`  GRPO-style baseline: group mean *including* the sample\n",
+    "* `E2`  leave-one-out (LOO) baseline: mean of the other `G-1`\n",
+    "* `E3`  tree: `G1` first-turn branches x `G2` second-turn siblings; per-turn LOO (turn-2 baseline = sibling mean, turn-1 baseline = other branches' mean value)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "feed6ed4",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:01:52.256933Z",
+     "iopub.status.busy": "2026-10-14T09:01:52.256766Z",
+     "iopub.status.idle": "2026-10-14T09:01:55.102052Z",
+     "shell.execute_reply": "2026-10-14T09:01:55.101235Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "E0     |E-gradJ|max=4.996e-16  trace Var=0.41416   (sum prob=1.000000000000)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "E0rtg  |E-gradJ|max=4.996e-16  trace Var=0.23686   (sum prob=1.000000000000)\n",
+      "E1     |E-gradJ|max=9.267e-02  trace Var=0.14155   (sum prob=1.000000000000)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "E2     |E-gradJ|max=8.327e-16  trace Var=0.25165   (sum prob=1.000000000000)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "BAD    |E-gradJ|max=7.857e-02  trace Var=0.19228   (sum prob=1.000000000000)\n",
+      "4. E0 (full-return, no baseline, G=4 episodes): unbiased                               err=4.996e-16  (tol 1e-15)\n",
+      "4. E0rtg (reward-to-go, G=4): unbiased                                                 err=4.996e-16  (tol 1e-15)\n",
+      "4. identity R_i - mean_all = (G-1)/G (R_i - mean_{-i}) on every group (max abs diff)   err=5.551e-17  (tol 1e-14)\n",
+      "4. GRPO mean baseline: E[g] = (G-1)/G grad J exactly                                   err=1.193e-15  (tol 1e-14)\n",
+      "5. leave-one-out baseline (episode-level, G=4): E[g] = grad J                          err=8.327e-16  (tol 1e-14)\n",
+      "own-outcome baseline: max |bias| = 0.07857014805316431\n",
+      "4. variance(E1 scaled by G/(G-1)) = variance(E2)                                       err=7.438e-15  (tol 1e-12)\n",
+      "tree LOO (2 branches x 2 siblings = 4 episodes): |E-gradJ|max=2.220e-16  trace Var=0.48614\n",
+      "6. tree two-level leave-one-out with per-turn baselines: E[g] = grad J                 err=2.220e-16  (tol 1e-14)\n",
+      "equal budget (4 episodes): trace Var  E0=0.4142  E0rtg=0.2369  flat LOO E2=0.2516  tree LOO=0.4861\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "offset c=0.0: trace Var  no-baseline=0.414  flat LOO=0.252  tree LOO=0.486\n",
+      "offset c=3.0: trace Var  no-baseline=4.908  flat LOO=0.252  tree LOO=0.486\n"
+     ]
+    }
+   ],
+   "source": [
+    "Gs = 4\n",
+    "leaf_ids = list(range(len(L)))\n",
+    "Pl = np.array([P[l] for l in L]); Rl = np.array([Rret[l] for l in L])\n",
+    "S1 = np.stack([s1(l[0]) for l in L]); S2 = np.stack([s2(*l) for l in L]); S12 = S1 + S2\n",
+    "R2l = np.array([r2[l] for l in L])\n",
+    "def run_flat(Gs, kind):\n",
+    "    m1 = np.zeros(nparam); m2 = 0.0; tot_prob = 0.0; maxdiff_1_2 = 0.0\n",
+    "    for grp in itertools.product(leaf_ids, repeat=Gs):\n",
+    "        pr = np.prod(Pl[list(grp)]); gi = list(grp)\n",
+    "        R = Rl[gi]\n",
+    "        if kind == 'E0':\n",
+    "            est = (R[:, None] * S12[gi]).mean(0)\n",
+    "        elif kind == 'E0rtg':\n",
+    "            est = (R[:, None] * S1[gi] + R2l[gi][:, None] * S2[gi]).mean(0)\n",
+    "        elif kind in ('E1', 'E2'):\n",
+    "            mean_all = R.mean(); loo = (R.sum() - R) / (Gs - 1)\n",
+    "            b = mean_all if kind == 'E1' else loo\n",
+    "            est = ((R - b)[:, None] * S12[gi]).mean(0)\n",
+    "            maxdiff_1_2 = max(maxdiff_1_2, float(np.abs((R - mean_all) - (Gs - 1) / Gs * (R - loo)).max()))\n",
+    "        elif kind == 'BAD':      # baseline depends on the sample's own turn-2 outcome: b_i = r2_i (for the turn-2 token), mean-LOO for turn 1\n",
+    "            loo = (R.sum() - R) / (Gs - 1)\n",
+    "            est = ((R - loo)[:, None] * S1[gi] + (R - R2l[gi])[:, None] * S2[gi]).mean(0)\n",
+    "        m1 += pr * est; m2 += pr * (est ** 2).sum(); tot_prob += pr\n",
+    "    return m1, m2 - (m1 ** 2).sum(), tot_prob, maxdiff_1_2\n",
+    "res = {}\n",
+    "for kind in ('E0', 'E0rtg', 'E1', 'E2', 'BAD'):\n",
+    "    res[kind] = run_flat(Gs, kind)\n",
+    "    print(f'{kind:6s} |E-gradJ|max={np.abs(res[kind][0]-gradJ).max():.3e}  trace Var={res[kind][1]:.5f}   (sum prob={res[kind][2]:.12f})')\n",
+    "check('4. E0 (full-return, no baseline, G=4 episodes): unbiased', np.abs(res['E0'][0] - gradJ).max(), 1e-15)\n",
+    "check('4. E0rtg (reward-to-go, G=4): unbiased', np.abs(res['E0rtg'][0] - gradJ).max(), 1e-15)\n",
+    "check('4. identity R_i - mean_all = (G-1)/G (R_i - mean_{-i}) on every group (max abs diff)', res['E1'][3], 1e-14)\n",
+    "check('4. GRPO mean baseline: E[g] = (G-1)/G grad J exactly', np.abs(res['E1'][0] - (Gs - 1) / Gs * gradJ).max(), 1e-14)\n",
+    "check('5. leave-one-out baseline (episode-level, G=4): E[g] = grad J', np.abs(res['E2'][0] - gradJ).max(), 1e-14)\n",
+    "bad_bias = np.abs(res['BAD'][0] - gradJ).max()\n",
+    "print('own-outcome baseline: max |bias| =', bad_bias)\n",
+    "assert bad_bias > 1e-3\n",
+    "CHECKS.append(('5. baseline depending on the sample own r2: |bias| (informational, must be > 0)', float(bad_bias), float('inf')))\n",
+    "# rescale E1 to compare variance\n",
+    "var_E1_scaled = res['E1'][1] * (Gs / (Gs - 1)) ** 2\n",
+    "check('4. variance(E1 scaled by G/(G-1)) = variance(E2)', abs(var_E1_scaled - res['E2'][1]), 1e-12)\n",
+    "\n",
+    "# --- tree estimator\n",
+    "G1, G2 = 2, 2\n",
+    "def run_tree():\n",
+    "    m1 = np.zeros(nparam); m2 = 0.0\n",
+    "    a_range = range(K1); b_range = range(K2)\n",
+    "    for a_s in itertools.product(a_range, repeat=G1):\n",
+    "        pa = np.prod([p1n[a] for a in a_s])\n",
+    "        for b_all in itertools.product(*[itertools.product(b_range, repeat=G2) for _ in range(G1)]):\n",
+    "            pb = np.prod([p2n[a, b] for a, bs in zip(a_s, b_all) for b in bs])\n",
+    "            Qhat = np.array([np.mean([Rret[(a, b)] for b in bs]) for a, bs in zip(a_s, b_all)])\n",
+    "            est = np.zeros(nparam)\n",
+    "            for i, (a, bs) in enumerate(zip(a_s, b_all)):\n",
+    "                b1 = (Qhat.sum() - Qhat[i]) / (G1 - 1)\n",
+    "                est += (Qhat[i] - b1) * s1(a) / G1\n",
+    "                Rs = np.array([Rret[(a, b)] for b in bs])\n",
+    "                for j, b in enumerate(bs):\n",
+    "                    b2 = (Rs.sum() - Rs[j]) / (G2 - 1)\n",
+    "                    est += (Rs[j] - b2) * s2(a, b) / (G1 * G2)\n",
+    "            pr = pa * pb; m1 += pr * est; m2 += pr * (est ** 2).sum()\n",
+    "    return m1, m2 - (m1 ** 2).sum()\n",
+    "m_tree, v_tree = run_tree()\n",
+    "print(f'tree LOO (2 branches x 2 siblings = 4 episodes): |E-gradJ|max={np.abs(m_tree-gradJ).max():.3e}  trace Var={v_tree:.5f}')\n",
+    "check('6. tree two-level leave-one-out with per-turn baselines: E[g] = grad J', np.abs(m_tree - gradJ).max(), 1e-14)\n",
+    "print(f'equal budget (4 episodes): trace Var  E0={res[\"E0\"][1]:.4f}  E0rtg={res[\"E0rtg\"][1]:.4f}  flat LOO E2={res[\"E2\"][1]:.4f}  tree LOO={v_tree:.4f}')\n",
+    "\n",
+    "# Value of baselines: add a constant c to every return (a \"reward offset\"; unbiasedness is unaffected, variance is not).\n",
+    "for c in (0.0, 3.0):\n",
+    "    Rl_save = Rl.copy(); Rret_save = dict(Rret)\n",
+    "    Rl = Rl_save + c; Rret = {k_: v_ + c for k_, v_ in Rret_save.items()}\n",
+    "    v0 = run_flat(Gs, 'E0')[1]; v2 = run_flat(Gs, 'E2')[1]; vt = run_tree()[1]\n",
+    "    print(f'offset c={c}: trace Var  no-baseline={v0:.3f}  flat LOO={v2:.3f}  tree LOO={vt:.3f}')\n",
+    "    Rl = Rl_save; Rret = Rret_save"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c2aa29c6",
+   "metadata": {},
+   "source": [
+    "## 7. Training only the first correction round on a stale context distribution (Goedel-Prover-V2-style RL data)\n",
+    "Use contexts `y1 ~ mu` (a frozen earlier policy) and train only `pi2`. Its expected gradient is\n",
+    "`sum_a mu(a) sum_b pi2(b|a) R(a,b) grad log pi2(b|a)` = `sum_a mu(a) grad_{theta2} Q(a)`. The full-episode gradient on the `theta2` block is\n",
+    "`sum_a pi1(a) grad_{theta2} Q(a)` and the full gradient has an additional `theta1` block. Compare."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "c9a05411",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:01:55.103908Z",
+     "iopub.status.busy": "2026-10-14T09:01:55.103731Z",
+     "iopub.status.idle": "2026-10-14T09:01:55.110368Z",
+     "shell.execute_reply": "2026-10-14T09:01:55.109451Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "7. truncated turn-2 gradient = sum_a mu(a) grad Q(a) (closed form)                     err=8.674e-18  (tol 1e-16)\n",
+      "theta2-block bias norm (context shift only): 0.10440110286572064   theta1-block (missing turn-1 term) norm: 0.4583795568241773\n",
+      "7. with mu = pi1 the theta2 block of the truncated gradient equals that of grad J      err=8.327e-17  (tol 1e-15)\n"
+     ]
+    }
+   ],
+   "source": [
+    "mu = np.array([0.6, 0.3, 0.1])\n",
+    "Qa = np.array([sum(p2n[a, b] * Rret[(a, b)] for b in range(K2)) for a in range(K1)])\n",
+    "E_trunc = np.zeros(nparam)\n",
+    "for a in range(K1):\n",
+    "    for b in range(K2):\n",
+    "        E_trunc += mu[a] * p2n[a, b] * Rret[(a, b)] * s2(a, b)\n",
+    "# closed form: sum_a mu(a) grad_{theta2} Q(a);  grad_{theta2[a]} Q(a) = pi2(b|a)(R(a,b) - Q(a)) over b\n",
+    "E_closed = np.zeros(nparam)\n",
+    "for a in range(K1):\n",
+    "    for b in range(K2):\n",
+    "        E_closed[K1 + a * K2 + b] = mu[a] * p2n[a, b] * (Rret[(a, b)] - Qa[a])\n",
+    "check('7. truncated turn-2 gradient = sum_a mu(a) grad Q(a) (closed form)', np.abs(E_trunc - E_closed).max(), 1e-16)\n",
+    "bias_ctx = E_trunc - gradJ\n",
+    "b2_block = bias_ctx[K1:]; b1_block = bias_ctx[:K1]\n",
+    "print('theta2-block bias norm (context shift only):', float(np.linalg.norm(b2_block)), '  theta1-block (missing turn-1 term) norm:', float(np.linalg.norm(b1_block)))\n",
+    "# if mu = pi1 the theta2 block is unbiased\n",
+    "E_same = np.zeros(nparam)\n",
+    "for a in range(K1):\n",
+    "    for b in range(K2):\n",
+    "        E_same += p1n[a] * p2n[a, b] * Rret[(a, b)] * s2(a, b)\n",
+    "check('7. with mu = pi1 the theta2 block of the truncated gradient equals that of grad J', np.abs(E_same[K1:] - gradJ[K1:]).max(), 1e-15)\n",
+    "CHECKS.append(('7. norm of the missing turn-1 block ||grad_theta1 J|| (informational, must be > 0)', float(np.linalg.norm(gradJ[:K1])), float('inf')))\n",
+    "assert np.linalg.norm(b1_block) > 1e-3 and np.linalg.norm(b2_block) > 1e-4"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a9073ea5",
+   "metadata": {},
+   "source": [
+    "## Verification table"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "b9986b19",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:01:55.112041Z",
+     "iopub.status.busy": "2026-10-14T09:01:55.111864Z",
+     "iopub.status.idle": "2026-10-14T09:01:55.115491Z",
+     "shell.execute_reply": "2026-10-14T09:01:55.114662Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "identity                                                                                    max abs err\n",
+      "1. E[R (score_turn1 + score_turn2)] = grad J (exact enumeration of 9 episodes)                8.327e-17\n",
+      "2. E[r1 * score_turn2] = 0                                                                    1.388e-17\n",
+      "2. reward-to-go estimator: E = grad J                                                         6.939e-17\n",
+      "3. norm of the bias from including feedback tokens in the loss (informational, must be > 0)    6.799e-01\n",
+      "4. E0 (full-return, no baseline, G=4 episodes): unbiased                                      4.996e-16\n",
+      "4. E0rtg (reward-to-go, G=4): unbiased                                                        4.996e-16\n",
+      "4. identity R_i - mean_all = (G-1)/G (R_i - mean_{-i}) on every group (max abs diff)          5.551e-17\n",
+      "4. GRPO mean baseline: E[g] = (G-1)/G grad J exactly                                          1.193e-15\n",
+      "5. leave-one-out baseline (episode-level, G=4): E[g] = grad J                                 8.327e-16\n",
+      "5. baseline depending on the sample own r2: |bias| (informational, must be > 0)               7.857e-02\n",
+      "4. variance(E1 scaled by G/(G-1)) = variance(E2)                                              7.438e-15\n",
+      "6. tree two-level leave-one-out with per-turn baselines: E[g] = grad J                        2.220e-16\n",
+      "7. truncated turn-2 gradient = sum_a mu(a) grad Q(a) (closed form)                            8.674e-18\n",
+      "7. with mu = pi1 the theta2 block of the truncated gradient equals that of grad J             8.327e-17\n",
+      "7. norm of the missing turn-1 block ||grad_theta1 J|| (informational, must be > 0)            4.584e-01\n",
+      "CHECKS_JSON:[[\"1. E[R (score_turn1 + score_turn2)] = grad J (exact enumeration of 9 episodes)\", 8.326672684688674e-17, 1e-15], [\"2. E[r1 * score_turn2] = 0\", 1.3877787807814457e-17, 1e-16], [\"2. reward-to-go estimator: E = grad J\", 6.938893903907228e-17, 1e-15], [\"3. norm of the bias from including feedback tokens in the loss (informational, must be > 0)\", 0.679933398028168, Infinity], [\"4. E0 (full-return, no baseline, G=4 episodes): unbiased\", 4.996003610813204e-16, 1e-15], [\"4. E0rtg (reward-to-go, G=4): unbiased\", 4.996003610813204e-16, 1e-15], [\"4. identity R_i - mean_all = (G-1)/G (R_i - mean_{-i}) on every group (max abs diff)\", 5.551115123125783e-17, 1e-14], [\"4. GRPO mean baseline: E[g] = (G-1)/G grad J exactly\", 1.1934897514720433e-15, 1e-14], [\"5. leave-one-out baseline (episode-level, G=4): E[g] = grad J\", 8.326672684688674e-16, 1e-14], [\"5. baseline depending on the sample own r2: |bias| (informational, must be > 0)\", 0.07857014805316431, Infinity], [\"4. variance(E1 scaled by G/(G-1)) = variance(E2)\", 7.438494264988549e-15, 1e-12], [\"6. tree two-level leave-one-out with per-turn baselines: E[g] = grad J\", 2.220446049250313e-16, 1e-14], [\"7. truncated turn-2 gradient = sum_a mu(a) grad Q(a) (closed form)\", 8.673617379884035e-18, 1e-16], [\"7. with mu = pi1 the theta2 block of the truncated gradient equals that of grad J\", 8.326672684688674e-17, 1e-15], [\"7. norm of the missing turn-1 block ||grad_theta1 J|| (informational, must be > 0)\", 0.4583795568241773, Infinity]]\n"
+     ]
+    }
+   ],
+   "source": [
+    "print()\n",
+    "print(f'{\"identity\":90s} {\"max abs err\":>12s}')\n",
+    "for n_, e_, t_ in CHECKS: print(f'{n_:90s} {e_:12.3e}')\n",
+    "print('CHECKS_JSON:' + json.dumps(CHECKS))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

code/notebooks/0205-model-averaging-weight-interpolation.ipynb

@@ -0,0 +1,746 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "5ac94761",
+   "metadata": {},
+   "source": [
+    "# 0205 -- Model averaging / weight interpolation `theta_alpha = (1-alpha) theta_0 + alpha theta`\n",
+    "\n",
+    "Convention (Goedel-Prover-V2 Sec. 2.3): `alpha` is the weight on the **fine-tuned** model; `alpha=0` is the base. (Their Sec. 3.6 text states\n",
+    "that alpha is \"the weight of the base model\" -- the opposite convention; see the note.)\n",
+    "\n",
+    "Identities / claims verified (fp64, CPU):\n",
+    "1. Logit-linear model: `softmax(z_alpha) = normalised pi_0^(1-alpha) pi_1^alpha` (geometric mixture), not the arithmetic mixture; naive (re-evaluate at theta_alpha) = efficient (interpolate logits).\n",
+    "2. Exponential tilt: `d/dalpha E_alpha[r] = Cov_alpha(g, r)`; for a binary tilt `p(alpha) = sigmoid(logit p0 + alpha c)` exactly.\n",
+    "3. Log-odds-linear first-order formulas: `d pass@k / d alpha = E_x[k p (1-p)^k Delta_x]`.\n",
+    "4. Curvature along the path for cross-entropy: `d^2 L/dalpha^2 = E_x Var_{y~softmax(z_alpha)}[Delta z_y]` (model-soups Eq. 1 with linear logits: soup = ensemble).\n",
+    "5. Loss along a line: `l(alpha) - chord = (1/2) alpha (alpha-1) d^T H d` for a quadratic; convex NLL has no barrier.\n",
+    "6. Linear mode connectivity: barrier for same-init fine-tunes vs independently trained nets (tiny MLP).\n",
+    "7. 30 random log-linear \"SFT -> RL\" pairs: pass@1 rises with alpha; pass@k (k=8,64) peaks at interior alpha."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "4ac1843a",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:49:30.138556Z",
+     "iopub.status.busy": "2026-10-14T09:49:30.138404Z",
+     "iopub.status.idle": "2026-10-14T09:49:31.926821Z",
+     "shell.execute_reply": "2026-10-14T09:49:31.925791Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n",
+    "import math, json\n",
+    "import numpy as np, torch\n",
+    "torch.set_default_dtype(torch.float64)\n",
+    "CHECKS = []\n",
+    "def check(name, err, tol):\n",
+    "    CHECKS.append((name, float(err), float(tol)))\n",
+    "    assert err <= tol, (name, float(err), tol)\n",
+    "    print(f'{name:86s} err={float(err):.3e}  (tol {tol:.0e})')\n",
+    "\n",
+    "g = np.random.RandomState(0)\n",
+    "m, K, d = 40, 12, 5\n",
+    "Phi = torch.tensor(g.randn(m, K, d)); V = torch.tensor((g.rand(m, K) < 0.12).astype(float))\n",
+    "V[V.sum(-1) == 0, 0] = 1.0\n",
+    "th0 = torch.tensor(g.randn(d) * 0.4); th1 = th0 + torch.tensor(g.randn(d) * 1.2)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "97ec1d44",
+   "metadata": {},
+   "source": [
+    "## 1. Geometric (log-linear) mixture\n",
+    "NAIVE: build `theta_alpha` and run the model. EFFICIENT: evaluate the two endpoint logit tables once and interpolate the logits.\n",
+    "Both equal the per-context normalised product `pi_0^(1-alpha) pi_1^alpha`. The arithmetic mixture `(1-alpha) pi_0 + alpha pi_1` (an ensemble) differs."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "91130ea5",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:49:31.929274Z",
+     "iopub.status.busy": "2026-10-14T09:49:31.929030Z",
+     "iopub.status.idle": "2026-10-14T09:49:31.936414Z",
+     "shell.execute_reply": "2026-10-14T09:49:31.935536Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1. naive (re-evaluate theta_alpha) vs efficient (interpolate logits)                   err=4.441e-16  (tol 1e-15)\n",
+      "1. softmax(z_alpha) vs normalised pi0^(1-a) pi1^a                                      err=4.441e-16  (tol 1e-15)\n",
+      "max |geometric - arithmetic mixture| for alpha in (0,1): 0.25554158391614934\n"
+     ]
+    }
+   ],
+   "source": [
+    "z0, z1 = Phi @ th0, Phi @ th1\n",
+    "pi0, pi1 = torch.softmax(z0, -1), torch.softmax(z1, -1)\n",
+    "worst_naive_eff = 0.0; worst_geo = 0.0; dist_arith = 0.0\n",
+    "for a in (0.0, 0.3, 0.7, 1.0, 1.25, -0.2):\n",
+    "    tha = (1 - a) * th0 + a * th1\n",
+    "    naive = torch.softmax(Phi @ tha, -1)\n",
+    "    eff = torch.softmax((1 - a) * z0 + a * z1, -1)\n",
+    "    geo = pi0 ** (1 - a) * pi1 ** a; geo = geo / geo.sum(-1, keepdim=True)\n",
+    "    ar = (1 - a) * pi0 + a * pi1\n",
+    "    worst_naive_eff = max(worst_naive_eff, float((naive - eff).abs().max()))\n",
+    "    worst_geo = max(worst_geo, float((naive - geo).abs().max()))\n",
+    "    if 0 < a < 1: dist_arith = max(dist_arith, float((naive - ar).abs().max()))\n",
+    "check('1. naive (re-evaluate theta_alpha) vs efficient (interpolate logits)', worst_naive_eff, 1e-15)\n",
+    "check('1. softmax(z_alpha) vs normalised pi0^(1-a) pi1^a', worst_geo, 1e-15)\n",
+    "print('max |geometric - arithmetic mixture| for alpha in (0,1):', dist_arith)\n",
+    "assert dist_arith > 1e-3"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "74e8c4e1",
+   "metadata": {},
+   "source": [
+    "## 2. Exponential tilt\n",
+    "If `pi_1 = pi_0 e^{g}/Z` then `pi_alpha = pi_0 e^{alpha g}/Z_alpha`, `d/dalpha E_alpha[r] = Cov_alpha(g, r)` (derivative of a log-partition function).\n",
+    "Binary tilt `g = c r`: odds multiply by `e^{alpha c}`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "1cc640d6",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:49:31.938137Z",
+     "iopub.status.busy": "2026-10-14T09:49:31.937942Z",
+     "iopub.status.idle": "2026-10-14T09:49:32.021196Z",
+     "shell.execute_reply": "2026-10-14T09:49:32.020228Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2. d/dalpha E_alpha[r] (autograd) = Cov_alpha(g, r)                                    err=5.204e-18  (tol 1e-16)\n",
+      "2. binary tilt: p(alpha) = sigmoid(logit p0 + alpha c)                                 err=0.000e+00  (tol 1e-15)\n"
+     ]
+    }
+   ],
+   "source": [
+    "gtilt = Phi @ (th1 - th0)                          # = z1 - z0, the tilt of each context\n",
+    "a = torch.tensor(0.4, requires_grad=True)\n",
+    "pa = torch.softmax(z0 + a * gtilt, -1)\n",
+    "Ea = (pa * V).sum(-1).mean(); Ea.backward()\n",
+    "pa_ = pa.detach(); cov = ((pa_ * gtilt * V).sum(-1) - (pa_ * gtilt).sum(-1) * (pa_ * V).sum(-1)).mean()\n",
+    "check('2. d/dalpha E_alpha[r] (autograd) = Cov_alpha(g, r)', abs(float(a.grad - cov)), 1e-16)\n",
+    "p0v = float((pi0 * V).sum(-1)[0]); c = 1.7\n",
+    "q = pi0[0] * torch.exp(c * V[0]); q = q / q.sum()\n",
+    "worst = 0.0\n",
+    "for al in (0.0, 0.25, 0.5, 1.0, 2.0):\n",
+    "    qa = pi0[0] * torch.exp(al * c * V[0]); qa = qa / qa.sum()\n",
+    "    p_al = float((qa * V[0]).sum()); logit = math.log(p0v / (1 - p0v)) + al * c\n",
+    "    worst = max(worst, abs(p_al - 1 / (1 + math.exp(-logit))))\n",
+    "check('2. binary tilt: p(alpha) = sigmoid(logit p0 + alpha c)', worst, 1e-15)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "87a33cf6",
+   "metadata": {},
+   "source": [
+    "## 2b. Scope of the covariance identity: single categorical vs autoregressive policy\n",
+    "For an autoregressive policy `pi_alpha(y) = prod_t softmax((1-alpha) z0 + alpha z1)(y_t | s_t)` the normaliser depends on the context, so\n",
+    "`d/dalpha log pi_alpha(y) = sum_t [g_t - E_{pi_alpha(.|s_t)} g]` (the *sequence score*), and `d/dalpha E_alpha[r] = Cov_alpha(score, r)`.\n",
+    "Replacing the score by `sum_t g_t` is NOT exact. Two-token toy, vocabulary 3, exact enumeration of the 9 sequences, central finite difference in `alpha`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "e54a128b",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:49:32.022698Z",
+     "iopub.status.busy": "2026-10-14T09:49:32.022530Z",
+     "iopub.status.idle": "2026-10-14T09:49:32.032914Z",
+     "shell.execute_reply": "2026-10-14T09:49:32.032036Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "finite difference dE/dalpha          : -0.19188850988727246\n",
+      "Cov_alpha(centred sequence score, r) : -0.19188850994087786\n",
+      "Cov_alpha(sum_t g_t, r)              : -0.17690360845304756\n",
+      "2b. autoregressive: dE/dalpha (finite diff.) = Cov_alpha(sequence score, r)            err=5.361e-11  (tol 1e-08)\n",
+      "2b. zero mean of the sequence score: E_alpha[score] = 0                                err=5.204e-17  (tol 1e-14)\n"
+     ]
+    }
+   ],
+   "source": [
+    "import itertools\n",
+    "rg_ar = np.random.RandomState(0); Vv_ = 3\n",
+    "z0_ = {(): rg_ar.randn(Vv_)}; z1_ = {(): rg_ar.randn(Vv_)}\n",
+    "for a_ in range(Vv_): z0_[(a_,)] = rg_ar.randn(Vv_); z1_[(a_,)] = rg_ar.randn(Vv_)\n",
+    "sm_ = lambda zz: np.exp(zz - zz.max()) / np.exp(zz - zz.max()).sum()\n",
+    "seqs_ = list(itertools.product(range(Vv_), repeat=2)); r_ar = rg_ar.rand(len(seqs_))\n",
+    "def dist_ar(al):\n",
+    "    return np.array([sm_((1 - al) * z0_[()] + al * z1_[()])[a_] * sm_((1 - al) * z0_[(a_,)] + al * z1_[(a_,)])[b_] for a_, b_ in seqs_])\n",
+    "al0 = 0.4; hh = 1e-6\n",
+    "fd = (dist_ar(al0 + hh) @ r_ar - dist_ar(al0 - hh) @ r_ar) / (2 * hh)\n",
+    "pdist = dist_ar(al0)\n",
+    "gtok = lambda ctx, y: (z1_[ctx] - z0_[ctx])[y]\n",
+    "cmean = lambda ctx: sm_((1 - al0) * z0_[ctx] + al0 * z1_[ctx]) @ (z1_[ctx] - z0_[ctx])\n",
+    "G_sum = np.array([gtok((), a_) + gtok((a_,), b_) for a_, b_ in seqs_])\n",
+    "G_score = np.array([(gtok((), a_) - cmean(())) + (gtok((a_,), b_) - cmean((a_,))) for a_, b_ in seqs_])\n",
+    "covf = lambda x: (pdist * x * r_ar).sum() - (pdist * x).sum() * (pdist * r_ar).sum()\n",
+    "print('finite difference dE/dalpha          :', fd)\n",
+    "print('Cov_alpha(centred sequence score, r) :', covf(G_score))\n",
+    "print('Cov_alpha(sum_t g_t, r)              :', covf(G_sum))\n",
+    "check('2b. autoregressive: dE/dalpha (finite diff.) = Cov_alpha(sequence score, r)', abs(fd - covf(G_score)), 1e-8)\n",
+    "assert abs(fd - covf(G_sum)) > 1e-3\n",
+    "CHECKS.append(('2b. autoregressive: error of the naive Cov_alpha(sum_t g_t, r) (informational, must be > 0)', float(abs(fd - covf(G_sum))), float('inf')))\n",
+    "check('2b. zero mean of the sequence score: E_alpha[score] = 0', abs((pdist * G_score).sum()), 1e-14)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e3649b2b",
+   "metadata": {},
+   "source": [
+    "## 3. First-order (log-odds-linear) analysis of pass@1 and pass@k\n",
+    "Assume `logit p_x(alpha) = l_x + alpha Delta_x`. Then `d p/dalpha = p(1-p) Delta`, so\n",
+    "`d pass@1/dalpha = E[p(1-p) Delta]`, `d pass@k/dalpha = E[k p (1-p)^k Delta]`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "0a878415",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:49:32.034304Z",
+     "iopub.status.busy": "2026-10-14T09:49:32.034149Z",
+     "iopub.status.idle": "2026-10-14T09:49:32.087078Z",
+     "shell.execute_reply": "2026-10-14T09:49:32.086240Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "3. d pass@k/dalpha = E[k p (1-p)^k Delta] vs autograd, k in {1,4,32}                   err=2.776e-17  (tol 1e-15)\n",
+      "k=  1: alpha* = 1.000  pass@k(0)=0.2776  pass@k(alpha*)=0.4355  pass@k(1)=0.4355\n",
+      "k=  8: alpha* = 0.525  pass@k(0)=0.6717  pass@k(alpha*)=0.6762  pass@k(1)=0.6738\n",
+      "k= 32: alpha* = 0.000  pass@k(0)=0.8998  pass@k(alpha*)=0.8998  pass@k(1)=0.7868\n",
+      "p maximising the marginal value of a log-odds shift, k p (1-p)^k : 1/(k+1) = {1: 0.5, 4: 0.2, 32: 0.030303030303030304}\n"
+     ]
+    }
+   ],
+   "source": [
+    "rg = np.random.RandomState(2)\n",
+    "# Two populations (synthetic, to exhibit the regimes): 65% \"learned\" prompts that the RL stage improves, 35% \"tail\" prompts (hard, p ~ 5%)\n",
+    "# that it degrades. logit p_x(alpha) = l_x + alpha Delta_x.\n",
+    "N = 2000; nt = int(0.35 * N); n1 = N - nt\n",
+    "l = torch.tensor(np.concatenate([rg.randn(n1) * 1.5 - 0.5, rg.randn(nt) * 0.5 - 3.0]))\n",
+    "Delta = torch.tensor(np.concatenate([rg.randn(n1) * 0.3 + 1.5, rg.randn(nt) * 0.3 - 1.2]))\n",
+    "def passk(al, k):\n",
+    "    p = torch.sigmoid(l + al * Delta); return (1 - (1 - p) ** k).mean()\n",
+    "worst = 0.0\n",
+    "for k in (1, 4, 32):\n",
+    "    for al in (0.0, 0.5, 1.0):\n",
+    "        a_ = torch.tensor(al, requires_grad=True)\n",
+    "        p = torch.sigmoid(l + a_ * Delta); (1 - (1 - p) ** k).mean().backward()\n",
+    "        p_ = p.detach(); formula = (k * p_ * (1 - p_) ** k * Delta).mean()\n",
+    "        worst = max(worst, abs(float(a_.grad - formula)))\n",
+    "check('3. d pass@k/dalpha = E[k p (1-p)^k Delta] vs autograd, k in {1,4,32}', worst, 1e-15)\n",
+    "als = np.linspace(0, 1, 201)\n",
+    "star = {}\n",
+    "for k in (1, 8, 32):\n",
+    "    vals = np.array([float(passk(a_, k)) for a_ in als]); star[k] = als[vals.argmax()]\n",
+    "    print(f'k={k:3d}: alpha* = {als[vals.argmax()]:.3f}  pass@k(0)={vals[0]:.4f}  pass@k(alpha*)={vals.max():.4f}  pass@k(1)={vals[-1]:.4f}')\n",
+    "assert star[1] == 1.0 and 0.0 < star[8] < 1.0 and star[32] < star[8]\n",
+    "CHECKS.append(('3. regimes: alpha*(k=1) = 1, 0 < alpha*(k=8) < 1, alpha*(k=32) < alpha*(k=8)  (alpha*: %s)' % {k_: round(float(v_), 3) for k_, v_ in star.items()}, 0.0, float('inf')))\n",
+    "peak_p = {k: 1.0 / (k + 1) for k in (1, 4, 32)}\n",
+    "print('p maximising the marginal value of a log-odds shift, k p (1-p)^k : 1/(k+1) =', peak_p)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6860011a",
+   "metadata": {},
+   "source": [
+    "## 3b. The end-point sign test is sufficient, not necessary, for an interior optimum\n",
+    "Two populations on logit-linear paths: half of the prompts `l=-3, Delta=-1` (declining), half `l=-8, Delta=+10` (late bloomers). `k=8`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "24bcee3d",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:49:32.088723Z",
+     "iopub.status.busy": "2026-10-14T09:49:32.088559Z",
+     "iopub.status.idle": "2026-10-14T09:49:32.094636Z",
+     "shell.execute_reply": "2026-10-14T09:49:32.093860Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "d/dalpha pass@8 at 0: -0.1152, at 1: -0.0622;  pass@8(0)=0.1624  pass@8(1)=0.5676  max=0.5792 at alpha=0.802\n",
+      "3b. derivative formula E[k p (1-p)^k Delta] vs finite difference (this population)     err=1.250e-10  (tol 1e-07)\n"
+     ]
+    }
+   ],
+   "source": [
+    "sig_ = lambda x: 1 / (1 + np.exp(-x))\n",
+    "pop_ = [(0.5, -3.0, -1.0), (0.5, -8.0, 10.0)]\n",
+    "pk_ = lambda a_, k_: sum(w_ * (1 - (1 - sig_(l_ + a_ * D_)) ** k_) for w_, l_, D_ in pop_)\n",
+    "dpk_ = lambda a_, k_: sum(w_ * k_ * sig_(l_ + a_ * D_) * (1 - sig_(l_ + a_ * D_)) ** k_ * D_ for w_, l_, D_ in pop_)\n",
+    "grid_ = np.linspace(0, 1, 2001); f_ = pk_(grid_, 8)\n",
+    "print(f'd/dalpha pass@8 at 0: {dpk_(0.0, 8):.4f}, at 1: {dpk_(1.0, 8):.4f};  pass@8(0)={f_[0]:.4f}  pass@8(1)={f_[-1]:.4f}  max={f_.max():.4f} at alpha={grid_[f_.argmax()]:.3f}')\n",
+    "assert dpk_(0.0, 8) < 0 and dpk_(1.0, 8) < 0 and 0 < f_.argmax() < len(grid_) - 1 and f_.max() > f_[-1] + 1e-3 and f_.max() > f_[0] + 1e-3\n",
+    "fd_ = (pk_(0.3 + 1e-6, 8) - pk_(0.3 - 1e-6, 8)) / 2e-6\n",
+    "check('3b. derivative formula E[k p (1-p)^k Delta] vs finite difference (this population)', abs(fd_ - dpk_(0.3, 8)), 1e-7)\n",
+    "CHECKS.append(('3b. interior global optimum with negative derivative at both ends (alpha* = %.3f, informational)' % grid_[f_.argmax()], float(grid_[f_.argmax()]), float('inf')))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b777f25a",
+   "metadata": {},
+   "source": [
+    "## 4. Curvature along the path: model-soups Eq. (1) for logit-linear models\n",
+    "Cross-entropy of a designated correct proof `y*` per context: `L(alpha) = -E_x log softmax(z_alpha)[y*]`. Then\n",
+    "`d^2 L / dalpha^2 = E_x Var_{y ~ softmax(z_alpha)}[Delta z_y]` (second derivative of a log-partition function) and the soup equals the logit ensemble."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "c4618972",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:49:32.095990Z",
+     "iopub.status.busy": "2026-10-14T09:49:32.095840Z",
+     "iopub.status.idle": "2026-10-14T09:49:32.101953Z",
+     "shell.execute_reply": "2026-10-14T09:49:32.101282Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "4. d^2 L/dalpha^2 (autograd) = E_x Var_{softmax(z_alpha)}[Delta z]                     err=4.441e-16  (tol 1e-15)\n",
+      "4. soup loss = logit-ensemble loss for logit-linear models                             err=0.000e+00  (tol 1e-16)\n"
+     ]
+    }
+   ],
+   "source": [
+    "ystar = torch.tensor(np.argmax(V.numpy(), axis=1))\n",
+    "def L_of(al): return -torch.log_softmax(z0 + al * (z1 - z0), -1)[torch.arange(m), ystar].mean()\n",
+    "al = torch.tensor(0.35, requires_grad=True)\n",
+    "L = L_of(al); g1, = torch.autograd.grad(L, al, create_graph=True); g2, = torch.autograd.grad(g1, al)\n",
+    "pz = torch.softmax(z0 + al.detach() * (z1 - z0), -1); dz = z1 - z0\n",
+    "var_dz = ((pz * dz ** 2).sum(-1) - (pz * dz).sum(-1) ** 2).mean()\n",
+    "check('4. d^2 L/dalpha^2 (autograd) = E_x Var_{softmax(z_alpha)}[Delta z]', abs(float(g2 - var_dz)), 1e-15)\n",
+    "# soup vs logit-ensemble loss difference (soups Eq.(1) predicts ~0 when logits are linear in theta)\n",
+    "ens_logits = (1 - al.detach()) * z0 + al.detach() * z1\n",
+    "L_ens = -torch.log_softmax(ens_logits, -1)[torch.arange(m), ystar].mean()\n",
+    "check('4. soup loss = logit-ensemble loss for logit-linear models', abs(float(L_of(al.detach()) - L_ens)), 1e-16)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e774bc82",
+   "metadata": {},
+   "source": [
+    "## 5. Loss along a line: chord identity and barriers\n",
+    "Quadratic `l(theta) = l_0 + g^T (theta-theta_0) + 1/2 (theta-theta_0)^T H (theta-theta_0)`, `d = theta_1-theta_0`, `kappa = d^T H d`:\n",
+    "`l(alpha) - [(1-alpha) l(0) + alpha l(1)] = (1/2) alpha (alpha-1) kappa`. For `H` PSD this is `<= 0` (no barrier); the averaged model beats the chord."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "a8781b64",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:49:32.103566Z",
+     "iopub.status.busy": "2026-10-14T09:49:32.103406Z",
+     "iopub.status.idle": "2026-10-14T09:49:32.127330Z",
+     "shell.execute_reply": "2026-10-14T09:49:32.126399Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "5. quadratic: l(alpha) - chord = alpha(alpha-1) kappa / 2                              err=9.770e-15  (tol 1e-12)\n",
+      "5. convex softmax-linear NLL: barrier max_alpha [l(alpha)-chord] <= 0                  err=0.000e+00  (tol 1e-14)\n",
+      "barrier of the convex NLL (should be <= 0): 0.0\n"
+     ]
+    }
+   ],
+   "source": [
+    "rq = np.random.RandomState(8)\n",
+    "A = rq.randn(7, 7); H = torch.tensor(A @ A.T); gv = torch.tensor(rq.randn(7)); dd = torch.tensor(rq.randn(7))\n",
+    "lq = lambda t: 0.3 + gv @ t + 0.5 * t @ H @ t\n",
+    "kappa = float(dd @ H @ dd)\n",
+    "worst = max(abs(float(lq(a_ * dd) - ((1 - a_) * lq(0 * dd) + a_ * lq(dd))) - 0.5 * a_ * (a_ - 1) * kappa) for a_ in np.linspace(-0.5, 1.5, 9))\n",
+    "check('5. quadratic: l(alpha) - chord = alpha(alpha-1) kappa / 2', worst, 1e-12)\n",
+    "# convex NLL of the log-linear policy: no barrier\n",
+    "def nll(al): return float(L_of(torch.tensor(al)))\n",
+    "grid = np.linspace(0, 1, 101)\n",
+    "bar = max(nll(a_) - ((1 - a_) * nll(0.0) + a_ * nll(1.0)) for a_ in grid)\n",
+    "check('5. convex softmax-linear NLL: barrier max_alpha [l(alpha)-chord] <= 0', max(0.0, bar), 1e-14)\n",
+    "print('barrier of the convex NLL (should be <= 0):', bar)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8f32e360",
+   "metadata": {},
+   "source": [
+    "## 6. Linear mode connectivity in a tiny non-convex net\n",
+    "2-D spiral, 2 classes. A 2-64-64-2 tanh MLP, full-batch Adam in fp64.\n",
+    "* Same-init: pretrain on task A, then two fine-tunes on a rotated/noisier task B with different seeds.\n",
+    "* Independent: two nets trained from different random inits on task B.\n",
+    "Barrier `B = max_alpha [ L(theta_alpha) - ((1-alpha) L(theta_0) + alpha L(theta_1)) ]` evaluated on the union of both tasks' data."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "94b56789",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:49:32.128822Z",
+     "iopub.status.busy": "2026-10-14T09:49:32.128648Z",
+     "iopub.status.idle": "2026-10-14T09:49:37.344398Z",
+     "shell.execute_reply": "2026-10-14T09:49:37.343352Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "same-init fine-tunes : endpoints L=0.001,0.001  barrier=+0.0000\n",
+      "independent inits    : endpoints L=0.000,0.000  barrier=+3.1884\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/vignesh/.opencode/tmp/ipykernel_606603/4015540266.py:25: UserWarning: Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.\n",
+      "Consider using tensor.detach() first. (Triggered internally at /pytorch/torch/csrc/autograd/generated/python_variable_methods.cpp:838.)\n",
+      "  L0, L1 = float(loss(P0, Xall, Yall)), float(loss(P1, Xall, Yall))\n"
+     ]
+    }
+   ],
+   "source": [
+    "def spiral(n, rot, noise, seed):\n",
+    "    r = np.random.RandomState(seed); t = np.linspace(0.2, 3.2, n)\n",
+    "    X = []; Y = []\n",
+    "    for c in (0, 1):\n",
+    "        ang = t * 2.0 + c * math.pi + rot\n",
+    "        X.append(np.stack([t * np.cos(ang), t * np.sin(ang)], 1) + r.randn(n, 2) * noise); Y.append(np.full(n, c))\n",
+    "    return torch.tensor(np.concatenate(X)), torch.tensor(np.concatenate(Y))\n",
+    "XA, YA = spiral(300, 0.0, 0.03, 1); XB, YB = spiral(300, 0.25, 0.08, 2)\n",
+    "Xall, Yall = torch.cat([XA, XB]), torch.cat([YA, YB])\n",
+    "def init(seed, h=64):\n",
+    "    r = torch.Generator().manual_seed(seed)\n",
+    "    P = [torch.randn(2, h, generator=r) / math.sqrt(2), torch.zeros(h), torch.randn(h, h, generator=r) / math.sqrt(h), torch.zeros(h),\n",
+    "         torch.randn(h, 2, generator=r) / math.sqrt(h), torch.zeros(2)]\n",
+    "    return [p.clone().requires_grad_(True) for p in P]\n",
+    "def fwd(P, X): return torch.tanh(torch.tanh(X @ P[0] + P[1]) @ P[2] + P[3]) @ P[4] + P[5]\n",
+    "def loss(P, X, Y): return torch.nn.functional.cross_entropy(fwd(P, X), Y)\n",
+    "def train(P, X, Y, steps, lr, seed=None, batch=None):\n",
+    "    o = torch.optim.Adam(P, lr=lr); gen = torch.Generator().manual_seed(seed if seed is not None else 0)\n",
+    "    for _ in range(steps):\n",
+    "        idx = torch.randperm(len(X), generator=gen)[:batch] if batch else slice(None)\n",
+    "        o.zero_grad(); loss(P, X[idx], Y[idx]).backward(); o.step()\n",
+    "    return [p.detach().clone().requires_grad_(True) for p in P]\n",
+    "def interp(P0, P1, a): return [((1 - a) * p0 + a * p1).detach() for p0, p1 in zip(P0, P1)]\n",
+    "def barrier(P0, P1):\n",
+    "    L0, L1 = float(loss(P0, Xall, Yall)), float(loss(P1, Xall, Yall))\n",
+    "    return max(float(loss(interp(P0, P1, a), Xall, Yall)) - ((1 - a) * L0 + a * L1) for a in np.linspace(0, 1, 41)), L0, L1\n",
+    "pre = train(init(0), XA, YA, 600, 1e-2)\n",
+    "ftA = train([p.detach().clone().requires_grad_(True) for p in pre], XB, YB, 300, 3e-3, seed=1, batch=64)\n",
+    "ftB = train([p.detach().clone().requires_grad_(True) for p in pre], XB, YB, 300, 3e-3, seed=2, batch=64)\n",
+    "indA = train(init(11), XB, YB, 900, 1e-2); indB = train(init(12), XB, YB, 900, 1e-2)\n",
+    "b_same, L0s, L1s = barrier(ftA, ftB); b_ind, L0i, L1i = barrier(indA, indB)\n",
+    "print(f'same-init fine-tunes : endpoints L={L0s:.3f},{L1s:.3f}  barrier={b_same:+.4f}')\n",
+    "print(f'independent inits    : endpoints L={L0i:.3f},{L1i:.3f}  barrier={b_ind:+.4f}')\n",
+    "assert b_ind > 0.1 and b_same < b_ind / 5\n",
+    "CHECKS.append(('6. barrier, independent inits (informational, large)', float(b_ind), float('inf')))\n",
+    "CHECKS.append(('6. barrier, same-init fine-tunes (informational, small)', float(b_same), float('inf')))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ac53ed4f",
+   "metadata": {},
+   "source": [
+    "## 7. SFT -> RL interpolation on 30 random log-linear policy families\n",
+    "Each family: 96 prompts, 24 candidate proofs, 6 features, ~8% of candidates verified. `theta_sft` = 80 steps of pass@1 ascent from a random init, `theta_rl` = 400 more.\n",
+    "For each family evaluate exact `pass@1, pass@8, pass@64` on `alpha in [0,1]` (101 points). NAIVE = rebuild `theta_alpha`; EFFICIENT = interpolate logits (asserted equal)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "1e316635",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:49:37.346390Z",
+     "iopub.status.busy": "2026-10-14T09:49:37.346111Z",
+     "iopub.status.idle": "2026-10-14T09:49:46.306787Z",
+     "shell.execute_reply": "2026-10-14T09:49:46.306034Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "families (of 30) with d pass@k/dalpha > 0 at alpha=0 / < 0 at alpha=1 : {1: [30, 0], 8: [14, 30], 64: [3, 30]}\n",
+      "families whose end-point derivative signs imply an interior optimum   : {8: 14, 64: 3}\n",
+      "7. naive (theta_alpha) vs efficient (logit interpolation), all families                err=6.661e-16  (tol 1e-15)\n",
+      "pass@1 non-decreasing in alpha: 30/30 families\n",
+      "interior argmax (0<alpha*<1): pass@8 14/30   pass@64 3/30\n",
+      "median alpha*: pass@8 0.00  pass@64 0.00\n",
+      "mean pass@1 : alpha=0 0.091  alpha=1 0.135 ;  at alpha*_64 0.091\n",
+      "mean pass@8 : alpha=0 0.417  alpha=1 0.294 ;  best 0.420\n",
+      "mean pass@64: alpha=0 0.779  alpha=1 0.466 ;  best 0.780\n",
+      "family 0 curves at alpha=0,0.25,0.5,0.75,1:\n",
+      "  pass@1   [0.0959 0.1068 0.1252 0.145  0.1595]\n",
+      "  pass@8   [0.4538 0.4447 0.4014 0.354  0.3184]\n",
+      "  pass@64  [0.8234 0.7851 0.691  0.5819 0.4962]\n"
+     ]
+    }
+   ],
+   "source": [
+    "def family(seed, m=96, K=24, d=6, dens=0.08, steps_sft=80, steps_rl=400, lr=0.5):\n",
+    "    r = np.random.RandomState(seed)\n",
+    "    Ph = torch.tensor(r.randn(m, K, d)); Vv = torch.tensor((r.rand(m, K) < dens).astype(float))\n",
+    "    t0 = torch.tensor(r.randn(d) * 0.3)\n",
+    "    def pof(th): return (torch.softmax(Ph @ th, -1) * Vv).sum(-1)\n",
+    "    th = t0.clone().requires_grad_(True); o = torch.optim.SGD([th], lr=lr)\n",
+    "    for _ in range(steps_sft): o.zero_grad(); (-pof(th).mean()).backward(); o.step()\n",
+    "    t_sft = th.detach().clone()\n",
+    "    for _ in range(steps_rl): o.zero_grad(); (-pof(th).mean()).backward(); o.step()\n",
+    "    return Ph, Vv, t_sft, th.detach().clone()\n",
+    "als = np.linspace(0, 1, 101)\n",
+    "rows = []; maxdiff = 0.0\n",
+    "for seed in range(30):\n",
+    "    Ph, Vv, ta, tb = family(seed)\n",
+    "    za, zb = Ph @ ta, Ph @ tb\n",
+    "    P1 = np.zeros(len(als)); P8 = np.zeros(len(als)); P64 = np.zeros(len(als))\n",
+    "    for i, a_ in enumerate(als):\n",
+    "        p_eff = (torch.softmax((1 - a_) * za + a_ * zb, -1) * Vv).sum(-1)\n",
+    "        if i % 25 == 0:\n",
+    "            p_naive = (torch.softmax(Ph @ ((1 - a_) * ta + a_ * tb), -1) * Vv).sum(-1)\n",
+    "            maxdiff = max(maxdiff, float((p_eff - p_naive).abs().max()))\n",
+    "        P1[i] = float(p_eff.mean()); P8[i] = float((1 - (1 - p_eff) ** 8).mean()); P64[i] = float((1 - (1 - p_eff) ** 64).mean())\n",
+    "    rows.append(dict(seed=seed, mono1=bool(np.all(np.diff(P1) >= -1e-12)), a8=als[P8.argmax()], a64=als[P64.argmax()],\n",
+    "                     p1_0=P1[0], p1_1=P1[-1], p8_0=P8[0], p8_1=P8[-1], p8_star=P8.max(), p64_0=P64[0], p64_1=P64[-1], p64_star=P64.max(),\n",
+    "                     p1_at64=P1[P64.argmax()]))\n",
+    "    if seed == 0: curves0 = (P1.copy(), P8.copy(), P64.copy())\n",
+    "# Exact end-point derivatives: d pass@k/dalpha at alpha = E_x[ k (1-p)^(k-1) Cov_alpha(g, r) ], g = z1 - z0 (no log-odds assumption needed).\n",
+    "def deriv_k(Ph, Vv, za, zb, a_, k):\n",
+    "    pz = torch.softmax((1 - a_) * za + a_ * zb, -1); gg = zb - za\n",
+    "    p = (pz * Vv).sum(-1); cov = (pz * gg * Vv).sum(-1) - (pz * gg).sum(-1) * p\n",
+    "    return float((k * (1 - p) ** (k - 1) * cov).mean())\n",
+    "sign_ok = 0; cnt = {1: [0, 0], 8: [0, 0], 64: [0, 0]}; inter_exact = {8: 0, 64: 0}\n",
+    "for seed in range(30):\n",
+    "    Ph, Vv, ta, tb = family(seed); za, zb = Ph @ ta, Ph @ tb\n",
+    "    for k in (1, 8, 64):\n",
+    "        d0, d1 = deriv_k(Ph, Vv, za, zb, 0.0, k), deriv_k(Ph, Vv, za, zb, 1.0, k)\n",
+    "        cnt[k][0] += d0 > 0; cnt[k][1] += d1 < 0\n",
+    "        if k in (8, 64):\n",
+    "            row = rows[seed]; astar = row['a8'] if k == 8 else row['a64']\n",
+    "            assert (d1 >= 0) or (astar < 1.0)          # derivative negative at alpha=1 => interpolating toward the base helps\n",
+    "            assert (d0 <= 0) or (astar > 0.0)          # derivative positive at alpha=0 => moving toward the fine-tune helps\n",
+    "            inter_exact[k] += (d0 > 0 and d1 < 0)\n",
+    "print('families (of 30) with d pass@k/dalpha > 0 at alpha=0 / < 0 at alpha=1 :', cnt)\n",
+    "print('families whose end-point derivative signs imply an interior optimum   :', inter_exact)\n",
+    "check('7. naive (theta_alpha) vs efficient (logit interpolation), all families', maxdiff, 1e-15)\n",
+    "import statistics as st\n",
+    "mono = sum(r_['mono1'] for r_ in rows)\n",
+    "inter8 = sum(0 < r_['a8'] < 1 for r_ in rows); inter64 = sum(0 < r_['a64'] < 1 for r_ in rows)\n",
+    "print(f'pass@1 non-decreasing in alpha: {mono}/30 families')\n",
+    "print(f'interior argmax (0<alpha*<1): pass@8 {inter8}/30   pass@64 {inter64}/30')\n",
+    "print('median alpha*: pass@8 %.2f  pass@64 %.2f' % (st.median(r_['a8'] for r_ in rows), st.median(r_['a64'] for r_ in rows)))\n",
+    "print('mean pass@1 : alpha=0 %.3f  alpha=1 %.3f ;  at alpha*_64 %.3f' % (st.mean(r_['p1_0'] for r_ in rows), st.mean(r_['p1_1'] for r_ in rows), st.mean(r_['p1_at64'] for r_ in rows)))\n",
+    "print('mean pass@8 : alpha=0 %.3f  alpha=1 %.3f ;  best %.3f' % (st.mean(r_['p8_0'] for r_ in rows), st.mean(r_['p8_1'] for r_ in rows), st.mean(r_['p8_star'] for r_ in rows)))\n",
+    "print('mean pass@64: alpha=0 %.3f  alpha=1 %.3f ;  best %.3f' % (st.mean(r_['p64_0'] for r_ in rows), st.mean(r_['p64_1'] for r_ in rows), st.mean(r_['p64_star'] for r_ in rows)))\n",
+    "print('family 0 curves at alpha=0,0.25,0.5,0.75,1:')\n",
+    "for nm, cv in zip(('pass@1', 'pass@8', 'pass@64'), curves0): print(f'  {nm:8s}', np.round(cv[[0, 25, 50, 75, 100]], 4))\n",
+    "assert mono == 30\n",
+    "CHECKS.append(('7. families with pass@1 monotone in alpha (count of 30, informational)', float(mono), float('inf')))\n",
+    "CHECKS.append(('7. families with interior argmax of pass@64 (count of 30, informational)', float(inter64), float('inf')))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "99bfedce",
+   "metadata": {},
+   "source": [
+    "## 8. Two practical traps: LoRA averaging and bf16 interpolation\n",
+    "(a) Averaging two LoRA adapters factor-wise is not averaging their weight updates: `((B1+B2)/2)((A1+A2)/2) = (1/4)(B1A1 + B2A2 + B1A2 + B2A1)`\n",
+    "whereas `(B1A1 + B2A2)/2` is what weight averaging means. Scaling a *merged* update `Delta W = BA` by alpha is exact.\n",
+    "(b) bf16 has an 8-bit significand (7 stored bits): unit roundoff 2^-8 ~ 3.9e-3, spacing between 2^-8 and 2^-7 of the weight. A small fine-tuning delta can vanish when stored/interpolated in bf16."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "6f52a593",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:49:46.308605Z",
+     "iopub.status.busy": "2026-10-14T09:49:46.308302Z",
+     "iopub.status.idle": "2026-10-14T09:49:46.524226Z",
+     "shell.execute_reply": "2026-10-14T09:49:46.523465Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "8a. factor-wise LoRA average minus weight average = (1/4)(B1A2+B2A1-B1A1-B2A2)         err=6.661e-16  (tol 1e-13)\n",
+      "||factor-wise avg - weight avg||_F = 14.505125959712405  relative to ||weight avg||_F: 0.6213290933983042\n",
+      "delta/weight =   1e-01:  fraction of weights whose bf16 value is unchanged by fine-tuning = 0.018;  bf16-arithmetic midpoint error / mean |delta| = 0.01\n",
+      "delta/weight =   1e-02:  fraction of weights whose bf16 value is unchanged by fine-tuning = 0.169;  bf16-arithmetic midpoint error / mean |delta| = 0.14\n",
+      "delta/weight =   1e-03:  fraction of weights whose bf16 value is unchanged by fine-tuning = 0.685;  bf16-arithmetic midpoint error / mean |delta| = 0.44\n",
+      "delta/weight =   1e-04:  fraction of weights whose bf16 value is unchanged by fine-tuning = 0.942;  bf16-arithmetic midpoint error / mean |delta| = 0.49\n"
+     ]
+    }
+   ],
+   "source": [
+    "rl = np.random.RandomState(1)\n",
+    "B1, A1, B2, A2 = [torch.tensor(rl.randn(16, 4)), torch.tensor(rl.randn(4, 16)), torch.tensor(rl.randn(16, 4)), torch.tensor(rl.randn(4, 16))]\n",
+    "fw = ((B1 + B2) / 2) @ ((A1 + A2) / 2); true_avg = (B1 @ A1 + B2 @ A2) / 2\n",
+    "cross = 0.25 * (B1 @ A2 + B2 @ A1) - 0.25 * (B1 @ A1 + B2 @ A2)     # fw - true_avg = (1/4)(B1A2+B2A1) - (1/4)(B1A1+B2A2)\n",
+    "check('8a. factor-wise LoRA average minus weight average = (1/4)(B1A2+B2A1-B1A1-B2A2)', (fw - true_avg - cross).abs().max(), 1e-13)\n",
+    "print('||factor-wise avg - weight avg||_F =', float((fw - true_avg).norm()), ' relative to ||weight avg||_F:', float((fw - true_avg).norm() / true_avg.norm()))\n",
+    "rb = torch.Generator().manual_seed(0)\n",
+    "w0 = (torch.randn(1_000_000, generator=rb) * 0.02)\n",
+    "for rel in (1e-1, 1e-2, 1e-3, 1e-4):\n",
+    "    dlt = torch.randn(1_000_000, generator=rb) * 0.02 * rel\n",
+    "    w1 = w0 + dlt\n",
+    "    w0b, w1b = w0.to(torch.bfloat16), w1.to(torch.bfloat16)\n",
+    "    lost = float((w0b == w1b).double().mean())\n",
+    "    half_bf = ((0.5 * w0b.double() + 0.5 * w1b.double()))                           # exact midpoint of the stored bf16 values\n",
+    "    half_bf16 = (0.5 * w0b + 0.5 * w1b).double()                                    # computed entirely in bf16\n",
+    "    err_bf16 = float((half_bf16 - half_bf).abs().mean() / (dlt.abs().mean()))\n",
+    "    print(f'delta/weight = {rel:7.0e}:  fraction of weights whose bf16 value is unchanged by fine-tuning = {lost:.3f};  bf16-arithmetic midpoint error / mean |delta| = {err_bf16:.2f}')\n",
+    "CHECKS.append(('8b. fraction of weights whose bf16 value is unchanged when delta/weight = 1e-3 (informational)', float((w0.to(torch.bfloat16) == (w0 + torch.randn(1_000_000, generator=rb) * 0.02 * 1e-3).to(torch.bfloat16)).double().mean()), float('inf')))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a3c399c8",
+   "metadata": {},
+   "source": [
+    "## Verification table"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "9a2c801c",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:49:46.525767Z",
+     "iopub.status.busy": "2026-10-14T09:49:46.525608Z",
+     "iopub.status.idle": "2026-10-14T09:49:46.529296Z",
+     "shell.execute_reply": "2026-10-14T09:49:46.528514Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "identity                                                                                    max abs err\n",
+      "1. naive (re-evaluate theta_alpha) vs efficient (interpolate logits)                          4.441e-16\n",
+      "1. softmax(z_alpha) vs normalised pi0^(1-a) pi1^a                                             4.441e-16\n",
+      "2. d/dalpha E_alpha[r] (autograd) = Cov_alpha(g, r)                                           5.204e-18\n",
+      "2. binary tilt: p(alpha) = sigmoid(logit p0 + alpha c)                                        0.000e+00\n",
+      "2b. autoregressive: dE/dalpha (finite diff.) = Cov_alpha(sequence score, r)                   5.361e-11\n",
+      "2b. autoregressive: error of the naive Cov_alpha(sum_t g_t, r) (informational, must be > 0)    1.498e-02\n",
+      "2b. zero mean of the sequence score: E_alpha[score] = 0                                       5.204e-17\n",
+      "3. d pass@k/dalpha = E[k p (1-p)^k Delta] vs autograd, k in {1,4,32}                          2.776e-17\n",
+      "3. regimes: alpha*(k=1) = 1, 0 < alpha*(k=8) < 1, alpha*(k=32) < alpha*(k=8)  (alpha*: {1: 1.0, 8: 0.525, 32: 0.0})    0.000e+00\n",
+      "3b. derivative formula E[k p (1-p)^k Delta] vs finite difference (this population)            1.250e-10\n",
+      "3b. interior global optimum with negative derivative at both ends (alpha* = 0.802, informational)    8.020e-01\n",
+      "4. d^2 L/dalpha^2 (autograd) = E_x Var_{softmax(z_alpha)}[Delta z]                            4.441e-16\n",
+      "4. soup loss = logit-ensemble loss for logit-linear models                                    0.000e+00\n",
+      "5. quadratic: l(alpha) - chord = alpha(alpha-1) kappa / 2                                     9.770e-15\n",
+      "5. convex softmax-linear NLL: barrier max_alpha [l(alpha)-chord] <= 0                         0.000e+00\n",
+      "6. barrier, independent inits (informational, large)                                          3.188e+00\n",
+      "6. barrier, same-init fine-tunes (informational, small)                                       0.000e+00\n",
+      "7. naive (theta_alpha) vs efficient (logit interpolation), all families                       6.661e-16\n",
+      "7. families with pass@1 monotone in alpha (count of 30, informational)                        3.000e+01\n",
+      "7. families with interior argmax of pass@64 (count of 30, informational)                      3.000e+00\n",
+      "8a. factor-wise LoRA average minus weight average = (1/4)(B1A2+B2A1-B1A1-B2A2)                6.661e-16\n",
+      "8b. fraction of weights whose bf16 value is unchanged when delta/weight = 1e-3 (informational)    6.853e-01\n",
+      "CHECKS_JSON:[[\"1. naive (re-evaluate theta_alpha) vs efficient (interpolate logits)\", 4.440892098500626e-16, 1e-15], [\"1. softmax(z_alpha) vs normalised pi0^(1-a) pi1^a\", 4.440892098500626e-16, 1e-15], [\"2. d/dalpha E_alpha[r] (autograd) = Cov_alpha(g, r)\", 5.204170427930421e-18, 1e-16], [\"2. binary tilt: p(alpha) = sigmoid(logit p0 + alpha c)\", 0.0, 1e-15], [\"2b. autoregressive: dE/dalpha (finite diff.) = Cov_alpha(sequence score, r)\", 5.3605397898337515e-11, 1e-08], [\"2b. autoregressive: error of the naive Cov_alpha(sum_t g_t, r) (informational, must be > 0)\", 0.014984901434224901, Infinity], [\"2b. zero mean of the sequence score: E_alpha[score] = 0\", 5.204170427930421e-17, 1e-14], [\"3. d pass@k/dalpha = E[k p (1-p)^k Delta] vs autograd, k in {1,4,32}\", 2.7755575615628914e-17, 1e-15], [\"3. regimes: alpha*(k=1) = 1, 0 < alpha*(k=8) < 1, alpha*(k=32) < alpha*(k=8)  (alpha*: {1: 1.0, 8: 0.525, 32: 0.0})\", 0.0, Infinity], [\"3b. derivative formula E[k p (1-p)^k Delta] vs finite difference (this population)\", 1.249767789257561e-10, 1e-07], [\"3b. interior global optimum with negative derivative at both ends (alpha* = 0.802, informational)\", 0.802, Infinity], [\"4. d^2 L/dalpha^2 (autograd) = E_x Var_{softmax(z_alpha)}[Delta z]\", 4.440892098500626e-16, 1e-15], [\"4. soup loss = logit-ensemble loss for logit-linear models\", 0.0, 1e-16], [\"5. quadratic: l(alpha) - chord = alpha(alpha-1) kappa / 2\", 9.769962616701378e-15, 1e-12], [\"5. convex softmax-linear NLL: barrier max_alpha [l(alpha)-chord] <= 0\", 0.0, 1e-14], [\"6. barrier, independent inits (informational, large)\", 3.1884148897650437, Infinity], [\"6. barrier, same-init fine-tunes (informational, small)\", 0.0, Infinity], [\"7. naive (theta_alpha) vs efficient (logit interpolation), all families\", 6.661338147750939e-16, 1e-15], [\"7. families with pass@1 monotone in alpha (count of 30, informational)\", 30.0, Infinity], [\"7. families with interior argmax of pass@64 (count of 30, informational)\", 3.0, Infinity], [\"8a. factor-wise LoRA average minus weight average = (1/4)(B1A2+B2A1-B1A1-B2A2)\", 6.661338147750939e-16, 1e-13], [\"8b. fraction of weights whose bf16 value is unchanged when delta/weight = 1e-3 (informational)\", 0.6853, Infinity]]\n"
+     ]
+    }
+   ],
+   "source": [
+    "print()\n",
+    "print(f'{\"identity\":90s} {\"max abs err\":>12s}')\n",
+    "for n_, e_, t_ in CHECKS: print(f'{n_:90s} {e_:12.3e}')\n",
+    "print('CHECKS_JSON:' + json.dumps(CHECKS))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

code/notebooks/0206-curriculum-pass-rate-band.ipynb

@@ -0,0 +1,545 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "2bf0eb1a",
+   "metadata": {},
+   "source": [
+    "# 0206 -- Curriculum by pass-rate band: signal of group advantages as a function of solve rate `p`\n",
+    "\n",
+    "Setting: one prompt with solve probability `p = sigmoid(z)`; `G` samples with binary rewards `r_i ~ Bern(p)`; `c = sum r_i`, `p_hat = c/G`.\n",
+    "Identities verified (fp64):\n",
+    "1. GRPO (unnormalised) surrogate gradient wrt the logit `z` equals `p_hat (1-p_hat)` for every `c` (autograd vs closed form).\n",
+    "2. `E[p_hat(1-p_hat)] = (1-1/G) p (1-p)`; the derivative of `p` wrt `z` is `p(1-p)`, maximal at `p=1/2` with value `1/4`.\n",
+    "3. Exact variance, SNR, and the dead-group probability `p^G + (1-p)^G`.\n",
+    "4. Std-normalised advantages: gradient `sqrt(p_hat(1-p_hat))`, expected value vs `sqrt(p(1-p))` (Dr. GRPO's reweighting).\n",
+    "5. General bound `||grad p||^2 <= p(1-p) E||s||^2` in a random log-linear policy family.\n",
+    "6. Pass-rate band `(0, tau]`: effective prompt weight `E[p_hat(1-p_hat) 1{0<p_hat<=tau}] / (p(1-p))`.\n",
+    "7. Expected informative fraction under Beta-distributed solve rates (closed form with Beta functions) and the oversampling cost of dynamic sampling."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "63f2be14",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:50:16.324611Z",
+     "iopub.status.busy": "2026-10-14T09:50:16.324416Z",
+     "iopub.status.idle": "2026-10-14T09:50:18.405260Z",
+     "shell.execute_reply": "2026-10-14T09:50:18.404171Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n",
+    "import math, json, itertools\n",
+    "from fractions import Fraction as Fr\n",
+    "from math import comb\n",
+    "import numpy as np, torch\n",
+    "from scipy.special import betaln\n",
+    "torch.set_default_dtype(torch.float64)\n",
+    "CHECKS = []\n",
+    "def check(name, err, tol):\n",
+    "    CHECKS.append((name, float(err), float(tol)))\n",
+    "    assert err <= tol, (name, float(err), tol)\n",
+    "    print(f'{name:90s} err={float(err):.3e}  (tol {tol:.0e})')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "535c098f",
+   "metadata": {},
+   "source": [
+    "## 1. The group-advantage gradient in the Bernoulli-logit model\n",
+    "`log pi(y_i) = r_i log p + (1-r_i) log(1-p)`, `p = sigmoid(z)`, so the score is `d log pi / dz = r_i - p`.\n",
+    "NAIVE: build the GRPO surrogate `-(1/G) sum_i A_i log pi(y_i)` with `A_i = r_i - p_hat` (detached) and differentiate with autograd, for every `c`.\n",
+    "EFFICIENT: `g_hat = (1/G) sum_i (r_i - p_hat)(r_i - p) = (1/G) sum_i (r_i - p_hat) r_i = p_hat (1-p_hat)` because `sum_i (r_i - p_hat) = 0` and `r_i^2 = r_i`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "fc678d9e",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:50:18.406857Z",
+     "iopub.status.busy": "2026-10-14T09:50:18.406631Z",
+     "iopub.status.idle": "2026-10-14T09:50:18.501912Z",
+     "shell.execute_reply": "2026-10-14T09:50:18.501228Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1. autograd surrogate gradient = p_hat (1 - p_hat), all c in 0..8, four p                  err=2.776e-17  (tol 1e-15)\n"
+     ]
+    }
+   ],
+   "source": [
+    "G = 8\n",
+    "worst = 0.0\n",
+    "for p in (0.03, 0.3, 0.5, 0.8):\n",
+    "    for c in range(G + 1):\n",
+    "        r = torch.tensor([1.0] * c + [0.0] * (G - c)); p_hat = r.mean()\n",
+    "        z = torch.tensor(math.log(p / (1 - p)), requires_grad=True)\n",
+    "        pz = torch.sigmoid(z)\n",
+    "        logpi = r * torch.log(pz) + (1 - r) * torch.log(1 - pz)\n",
+    "        loss = -((r - p_hat).detach() * logpi).mean(); loss.backward()\n",
+    "        g_auto = -z.grad                                  # ascent direction\n",
+    "        worst = max(worst, abs(float(g_auto - p_hat * (1 - p_hat))))\n",
+    "check('1. autograd surrogate gradient = p_hat (1 - p_hat), all c in 0..8, four p', worst, 1e-15)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3ec78669",
+   "metadata": {},
+   "source": [
+    "## 2. Expectation: `(1-1/G) p (1-p)`\n",
+    "`E[p_hat(1-p_hat)] = (E p_hat - E p_hat^2) = p - (p(1-p)/G + p^2) = p(1-p)(1 - 1/G)`. Checked with exact rationals."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "609db332",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:50:18.503413Z",
+     "iopub.status.busy": "2026-10-14T09:50:18.503251Z",
+     "iopub.status.idle": "2026-10-14T09:50:18.514651Z",
+     "shell.execute_reply": "2026-10-14T09:50:18.513962Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2. E[p_hat(1-p_hat)] = (1-1/G) p(1-p), exact rationals                                     err=0.000e+00  (tol 0e+00)\n",
+      "argmax_p p(1-p) = 0.5  max = 0.25\n",
+      "2. max of p(1-p) is 1/4 at p = 1/2                                                         err=0.000e+00  (tol 1e-12)\n"
+     ]
+    }
+   ],
+   "source": [
+    "def E_exact(G, p): \n",
+    "    p = Fr(p); return sum(comb(G, c) * p ** c * (1 - p) ** (G - c) * Fr(c, G) * (1 - Fr(c, G)) for c in range(G + 1))\n",
+    "worst = 0.0\n",
+    "for G_ in (2, 4, 8, 16, 64):\n",
+    "    for p in (Fr(1, 50), Fr(1, 5), Fr(1, 2), Fr(9, 10)):\n",
+    "        worst = max(worst, float(abs(E_exact(G_, p) - (1 - Fr(1, G_)) * p * (1 - p))))\n",
+    "check('2. E[p_hat(1-p_hat)] = (1-1/G) p(1-p), exact rationals', worst, 0.0)\n",
+    "ps = np.linspace(0, 1, 100001); sig = ps * (1 - ps)\n",
+    "print('argmax_p p(1-p) =', ps[sig.argmax()], ' max =', sig.max())\n",
+    "check('2. max of p(1-p) is 1/4 at p = 1/2', abs(sig.max() - 0.25) + abs(ps[sig.argmax()] - 0.5), 1e-12)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8abe5ba6",
+   "metadata": {},
+   "source": [
+    "## 3. Variance, SNR, dead groups\n",
+    "`Var[p_hat(1-p_hat)]` exactly (binomial sum). SNR = `E / sd`. Dead group: `c in {0, G}` with probability `p^G + (1-p)^G`; there `g_hat = 0` exactly."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "530a2701",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:50:18.515964Z",
+     "iopub.status.busy": "2026-10-14T09:50:18.515809Z",
+     "iopub.status.idle": "2026-10-14T09:50:18.605919Z",
+     "shell.execute_reply": "2026-10-14T09:50:18.605072Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "G=8   p     E[g]     sd[g]    SNR     dead prob   informative fraction\n",
+      "      0.01  0.00866  0.03021   0.287  0.92274   0.07726\n",
+      "      0.05  0.04156  0.06119   0.679  0.66342   0.33658\n",
+      "      0.10  0.07875  0.07572   1.040  0.43047   0.56953\n",
+      "      0.25  0.16406  0.07381   2.223  0.10013   0.89987\n",
+      "      0.50  0.21875  0.04134   5.292  0.00781   0.99219\n",
+      "      0.75  0.16406  0.07381   2.223  0.10013   0.89987\n",
+      "      0.90  0.07875  0.07572   1.040  0.43047   0.56953\n",
+      "      0.99  0.00866  0.03021   0.287  0.92274   0.07726\n",
+      "3. MC mean of p_hat(1-p_hat) vs exact (in s.e.)                                            err=7.820e-01  (tol 4e+00)\n",
+      "3. MC dead-group fraction vs p^G + (1-p)^G (in s.e.)                                       err=3.515e-01  (tol 4e+00)\n",
+      "E[g] is largest at p = 0.5 ; SNR is largest at p = 0.5\n",
+      "dead prob at p=0.5, G=8: 0.0078125  at p=0.9: 0.4304672200000001\n"
+     ]
+    }
+   ],
+   "source": [
+    "def moments(G, p, f):\n",
+    "    pm = np.array([comb(G, c) * p ** c * (1 - p) ** (G - c) for c in range(G + 1)])\n",
+    "    v = np.array([f(c / G) for c in range(G + 1)]); m1 = (pm * v).sum(); m2 = (pm * v ** 2).sum()\n",
+    "    return m1, m2 - m1 ** 2\n",
+    "f_un = lambda ph: ph * (1 - ph)\n",
+    "print('G=8   p     E[g]     sd[g]    SNR     dead prob   informative fraction')\n",
+    "tab = []\n",
+    "for p in (0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.99):\n",
+    "    m, v = moments(8, p, f_un); dead = p ** 8 + (1 - p) ** 8\n",
+    "    tab.append((p, m, math.sqrt(v), m / math.sqrt(v) if v > 0 else float('nan'), dead))\n",
+    "    print(f'     {p:5.2f}  {m:.5f}  {math.sqrt(v):.5f}  {tab[-1][3]:6.3f}  {dead:.5f}   {1 - dead:.5f}')\n",
+    "# MC check of the exact moments\n",
+    "rgn = np.random.RandomState(0); p = 0.2; cs = rgn.binomial(8, p, size=2_000_000); ph = cs / 8; gh = ph * (1 - ph)\n",
+    "m, v = moments(8, p, f_un)\n",
+    "check('3. MC mean of p_hat(1-p_hat) vs exact (in s.e.)', abs(gh.mean() - m) / (gh.std() / math.sqrt(len(gh))), 4.0)\n",
+    "check('3. MC dead-group fraction vs p^G + (1-p)^G (in s.e.)', abs(((cs == 0) | (cs == 8)).mean() - (p ** 8 + (1 - p) ** 8)) / math.sqrt(0.5 * 0.5 / len(cs)), 4.0)\n",
+    "sig_peak = max(tab, key=lambda t: t[1]); snr_peak = max(tab, key=lambda t: t[3] if t[3] == t[3] else -1)\n",
+    "print('E[g] is largest at p =', sig_peak[0], '; SNR is largest at p =', snr_peak[0])\n",
+    "dead_at_half = 0.5 ** 8 * 2\n",
+    "print('dead prob at p=0.5, G=8:', dead_at_half, ' at p=0.9:', 0.9 ** 8 + 0.1 ** 8)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a0e51cdd",
+   "metadata": {},
+   "source": [
+    "## 4. Std-normalised advantages (GRPO): gradient `sqrt(p_hat (1 - p_hat))`\n",
+    "`A_i = (r_i - p_hat)/sigma_hat`, `sigma_hat = sqrt(p_hat(1-p_hat))` (population std). Then `g_hat = (1/G) sum_i A_i (r_i - p) = p_hat(1-p_hat)/sigma_hat = sigma_hat`\n",
+    "(for `sigma_hat > 0`; zero otherwise). Compare `E[sigma_hat]` with `sqrt(p(1-p))` and the *relative* weight of a prompt versus the unnormalised estimator."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "0a8dd37b",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:50:18.607353Z",
+     "iopub.status.busy": "2026-10-14T09:50:18.607193Z",
+     "iopub.status.idle": "2026-10-14T09:50:18.617653Z",
+     "shell.execute_reply": "2026-10-14T09:50:18.616857Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "4. std-normalised group gradient = sigma_hat (direct sum vs closed form)                   err=1.110e-16  (tol 1e-15)\n",
+      "4. unbiased-std normalisation: group gradient = sigma_hat sqrt((G-1)/G)                    err=1.110e-16  (tol 1e-15)\n",
+      "p      E[sigma_hat]  sqrt(p(1-p))   weight ratio  E[sigma_hat]/E[p_hat(1-p_hat)]   (unnormalised = 1 by definition)\n",
+      "0.02   0.05043      0.14000         2.941\n",
+      "0.10   0.20949      0.30000         2.660\n",
+      "0.25   0.37990      0.43301         2.316\n",
+      "0.50   0.46391      0.50000         2.121\n",
+      "0.90   0.20949      0.30000         2.660\n",
+      "relative prompt weight p=0.02 vs p=0.5: unnormalised 0.0784, std-normalised 0.1087\n",
+      "informative group c=1 vs c=4 (G=8): unnormalised 0.438, std-normalised 0.661, ratio of the two = 1/sigma_hat ratio = 1.512\n",
+      "4. std-normalisation re-weights c=1 vs c=4 groups by exactly sigma(4)/sigma(1)             err=2.220e-16  (tol 1e-14)\n"
+     ]
+    }
+   ],
+   "source": [
+    "worst = 0.0\n",
+    "for p in (0.1, 0.5):\n",
+    "    for c in range(1, G):\n",
+    "        r = torch.tensor([1.0] * c + [0.0] * (G - c)); ph = r.mean(); sg = torch.sqrt(ph * (1 - ph))\n",
+    "        A = (r - ph) / sg\n",
+    "        g_direct = float((A * (r - p)).mean())\n",
+    "        worst = max(worst, abs(g_direct - float(sg)))\n",
+    "check('4. std-normalised group gradient = sigma_hat (direct sum vs closed form)', worst, 1e-15)\n",
+    "# Same with the *unbiased* sample std (torch.std, G-1 denominator) used by many implementations: g = sigma_hat * sqrt((G-1)/G)\n",
+    "worst_u = 0.0\n",
+    "for p in (0.1, 0.5):\n",
+    "    for c in range(1, G):\n",
+    "        r = torch.tensor([1.0] * c + [0.0] * (G - c)); ph = r.mean(); s_unb = r.std(unbiased=True)\n",
+    "        g_direct = float((((r - ph) / s_unb) * (r - p)).mean())\n",
+    "        worst_u = max(worst_u, abs(g_direct - float(torch.sqrt(ph * (1 - ph))) * math.sqrt((G - 1) / G)))\n",
+    "check('4. unbiased-std normalisation: group gradient = sigma_hat sqrt((G-1)/G)', worst_u, 1e-15)\n",
+    "print('p      E[sigma_hat]  sqrt(p(1-p))   weight ratio  E[sigma_hat]/E[p_hat(1-p_hat)]   (unnormalised = 1 by definition)')\n",
+    "for p in (0.02, 0.1, 0.25, 0.5, 0.9):\n",
+    "    m_n, _ = moments(8, p, lambda ph: math.sqrt(ph * (1 - ph))); m_u, _ = moments(8, p, f_un)\n",
+    "    print(f'{p:4.2f}   {m_n:.5f}      {math.sqrt(p*(1-p)):.5f}      {m_n/m_u:8.3f}')\n",
+    "m_n05, _ = moments(8, 0.5, lambda ph: math.sqrt(ph * (1 - ph))); m_n02, _ = moments(8, 0.02, lambda ph: math.sqrt(ph * (1 - ph)))\n",
+    "m_u05, _ = moments(8, 0.5, f_un); m_u02, _ = moments(8, 0.02, f_un)\n",
+    "print('relative prompt weight p=0.02 vs p=0.5: unnormalised %.4f, std-normalised %.4f' % (m_u02 / m_u05, m_n02 / m_n05))\n",
+    "assert m_n02 / m_n05 > m_u02 / m_u05\n",
+    "# conditional on being informative, a group with c=1 of G=8 vs c=4 of 8: relative gradient magnitudes\n",
+    "sg = lambda c, G_=8: math.sqrt((c / G_) * (1 - c / G_)); un = lambda c, G_=8: (c / G_) * (1 - c / G_)\n",
+    "print('informative group c=1 vs c=4 (G=8): unnormalised %.3f, std-normalised %.3f, ratio of the two = 1/sigma_hat ratio = %.3f' % (un(1) / un(4), sg(1) / sg(4), (sg(1) / sg(4)) / (un(1) / un(4))))\n",
+    "check('4. std-normalisation re-weights c=1 vs c=4 groups by exactly sigma(4)/sigma(1)', abs((sg(1) / sg(4)) / (un(1) / un(4)) - sg(4) / sg(1)), 1e-14)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bd6aa13f",
+   "metadata": {},
+   "source": [
+    "## 5. General bound: `||grad p||^2 <= p(1-p) E||s||^2` (Cauchy-Schwarz), and `grad p = 0` at `p in {0,1}`\n",
+    "`grad p = E[(r - p) s]` (baseline identity), so `||grad p|| <= sqrt(E(r-p)^2) sqrt(E||s||^2) = sqrt(p(1-p)) sqrt(E||s||^2)`.\n",
+    "Verified on 300 prompts of a log-linear policy family (K=16 outputs, d=6 features) with exact `grad p` from autograd."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "fbf95af9",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:50:18.618988Z",
+     "iopub.status.busy": "2026-10-14T09:50:18.618837Z",
+     "iopub.status.idle": "2026-10-14T09:50:18.687537Z",
+     "shell.execute_reply": "2026-10-14T09:50:18.686732Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "max ratio ||grad p||^2 / [p(1-p) E||s||^2] = 0.3109090719429734   (must be <= 1)\n",
+      "5. Cauchy-Schwarz bound ||grad p||^2 <= p(1-p) E||s||^2 (max violation of ratio <= 1)      err=0.000e+00  (tol 0e+00)\n",
+      "  p in [0,0.001): n= 28  mean ||grad p|| = 0.0000\n",
+      "  p in [0.001,0.05): n= 24  mean ||grad p|| = 0.0791\n",
+      "  p in [0.05,0.3): n=123  mean ||grad p|| = 0.2233\n",
+      "  p in [0.3,0.7): n=112  mean ||grad p|| = 0.3436\n",
+      "  p in [0.7,0.95): n= 13  mean ||grad p|| = 0.3216\n",
+      "p=0 prompt  ||grad p|| = 0.0\n",
+      "p=1 prompt  ||grad p|| = 0.0\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/vignesh/.opencode/tmp/ipykernel_607906/3044780012.py:10: UserWarning: Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.\n",
+      "Consider using tensor.detach() first. (Triggered internally at /pytorch/torch/csrc/autograd/generated/python_variable_methods.cpp:838.)\n",
+      "  Es2 = float((pi.detach() * (S_ ** 2).sum(1)).sum()); pf = float(p)\n"
+     ]
+    }
+   ],
+   "source": [
+    "rg2 = np.random.RandomState(3); m_, K_, d_ = 300, 16, 6\n",
+    "Ph = torch.tensor(rg2.randn(m_, K_, d_)); Vv = torch.tensor((rg2.rand(m_, K_) < rg2.rand(m_, 1) * 0.6).astype(float))\n",
+    "th = torch.tensor(rg2.randn(d_) * 0.5, requires_grad=True)\n",
+    "ratios = []; grad_norms = []; pvals = []\n",
+    "for x in range(m_):\n",
+    "    th.grad = None\n",
+    "    pi = torch.softmax(Ph[x] @ th, -1); p = (pi * Vv[x]).sum(); p.backward()\n",
+    "    gp = th.grad.clone()\n",
+    "    S_ = Ph[x] - (pi.detach()[:, None] * Ph[x]).sum(0, keepdim=True)          # score of each output y wrt theta: phi(x,y) - E phi\n",
+    "    Es2 = float((pi.detach() * (S_ ** 2).sum(1)).sum()); pf = float(p)\n",
+    "    bound = pf * (1 - pf) * Es2\n",
+    "    ratios.append(float(gp @ gp) / bound if bound > 0 else 0.0); grad_norms.append(float(gp.norm())); pvals.append(pf)\n",
+    "print('max ratio ||grad p||^2 / [p(1-p) E||s||^2] =', max(ratios), '  (must be <= 1)')\n",
+    "check('5. Cauchy-Schwarz bound ||grad p||^2 <= p(1-p) E||s||^2 (max violation of ratio <= 1)', max(0.0, max(ratios) - 1.0), 0.0)\n",
+    "pv = np.array(pvals); gn = np.array(grad_norms)\n",
+    "for lo, hi in [(0, 0.001), (0.001, 0.05), (0.05, 0.3), (0.3, 0.7), (0.7, 0.95), (0.95, 1.0001)]:\n",
+    "    sel = (pv >= lo) & (pv < hi)\n",
+    "    if sel.sum(): print(f'  p in [{lo},{hi}): n={int(sel.sum()):3d}  mean ||grad p|| = {gn[sel].mean():.4f}')\n",
+    "Vz = Vv.clone(); Vz[0, :] = 0.0; Vo = Vv.clone(); Vo[1, :] = 1.0\n",
+    "for name_, Vt in (('p=0 prompt', Vz[0]), ('p=1 prompt', Vo[1])):\n",
+    "    th.grad = None; pi = torch.softmax(Ph[0] @ th, -1); (pi * Vt).sum().backward()\n",
+    "    print(name_, ' ||grad p|| =', float(th.grad.norm()))\n",
+    "    assert float(th.grad.norm()) < 1e-15"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ca8a07b9",
+   "metadata": {},
+   "source": [
+    "## 6. Pass-rate band `(0, tau]` as deliberate re-weighting\n",
+    "Keep a group iff `0 < p_hat <= tau`. Its expected gradient contribution is `E[p_hat(1-p_hat) 1{0 < p_hat <= tau}]` versus `p(1-p)` for the true derivative.\n",
+    "The ratio `w_tau(p)` is the effective weight multiplier a prompt with true solve rate `p` receives."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "bb33c172",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:50:18.689048Z",
+     "iopub.status.busy": "2026-10-14T09:50:18.688891Z",
+     "iopub.status.idle": "2026-10-14T09:50:18.820659Z",
+     "shell.execute_reply": "2026-10-14T09:50:18.819967Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "6. band weight w_tau(p): closed form vs simulation (p=.7,G=8,tau=.75)                      err=3.628e-04  (tol 3e-03)\n",
+      "effective weight multiplier w_tau(p) (1 - 1/G = no band):\n",
+      "   p      G=8,tau=0.75   G=16,tau=0.75   G=8,tau=1 (DAPO: 0<p_hat<1)\n",
+      "  0.02     0.8750        0.9375        0.8750\n",
+      "  0.10     0.8750        0.9375        0.8750\n",
+      "  0.30     0.8744        0.9375        0.8750\n",
+      "  0.50     0.8613        0.9314        0.8750\n",
+      "  0.70     0.7721        0.7867        0.8750\n",
+      "  0.80     0.6456        0.5175        0.8750\n",
+      "  0.90     0.4100        0.1485        0.8750\n",
+      "  0.97     0.1461        0.0072        0.8750\n",
+      "6. no band (tau=1) gives exactly 1-1/G for all p                                           err=3.331e-16  (tol 1e-14)\n"
+     ]
+    }
+   ],
+   "source": [
+    "def w_band(G, p, tau):\n",
+    "    pm = np.array([comb(G, c) * p ** c * (1 - p) ** (G - c) for c in range(G + 1)])\n",
+    "    keep = np.array([1.0 if (0 < c / G <= tau) else 0.0 for c in range(G + 1)])\n",
+    "    val = np.array([(c / G) * (1 - c / G) for c in range(G + 1)])\n",
+    "    return (pm * keep * val).sum() / (p * (1 - p))\n",
+    "# direct simulation check\n",
+    "rs = np.random.RandomState(1); p = 0.7; G_ = 8; tau = 0.75\n",
+    "cs = rs.binomial(G_, p, size=3_000_000); ph = cs / G_\n",
+    "sim = (ph * (1 - ph) * ((ph > 0) & (ph <= tau))).mean() / (p * (1 - p))\n",
+    "check('6. band weight w_tau(p): closed form vs simulation (p=.7,G=8,tau=.75)', abs(sim - w_band(G_, p, tau)), 3e-3)\n",
+    "print('effective weight multiplier w_tau(p) (1 - 1/G = no band):')\n",
+    "print('   p      G=8,tau=0.75   G=16,tau=0.75   G=8,tau=1 (DAPO: 0<p_hat<1)')\n",
+    "for p in (0.02, 0.1, 0.3, 0.5, 0.7, 0.8, 0.9, 0.97):\n",
+    "    print(f'  {p:4.2f}     {w_band(8,p,0.75):.4f}        {w_band(16,p,0.75):.4f}        {w_band(8,p,1.0):.4f}')\n",
+    "check('6. no band (tau=1) gives exactly 1-1/G for all p', max(abs(w_band(8, p, 1.0) - (1 - 1 / 8)) for p in (0.02, 0.3, 0.9)), 1e-14)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4e142cd1",
+   "metadata": {},
+   "source": [
+    "## 7. Informative fraction and the cost of dynamic sampling\n",
+    "Population of prompts with `p ~ Beta(a,b)`. `E[p^G] = B(a+G,b)/B(a,b)`, `E[(1-p)^G] = B(a,b+G)/B(a,b)`. Informative fraction `f = 1 - E[p^G] - E[(1-p)^G]`.\n",
+    "To fill a batch of `B` informative groups one must draw `B/f` prompts in expectation (and generate `G B / f` samples)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "f49460fb",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:50:18.822182Z",
+     "iopub.status.busy": "2026-10-14T09:50:18.822029Z",
+     "iopub.status.idle": "2026-10-14T09:50:21.136390Z",
+     "shell.execute_reply": "2026-10-14T09:50:21.135678Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "7. informative fraction: Beta-function closed form vs MC (max z)                           err=9.161e-01  (tol 4e+00)\n",
+      "population            G    informative fraction f   oversampling factor 1/f\n",
+      "hard-heavy Beta(0.5,1)     8     0.6416                   1.559\n",
+      "hard-heavy Beta(0.5,1)    16     0.7532                   1.328\n",
+      "hard-heavy Beta(0.5,1)    32     0.8298                   1.205\n",
+      "uniform Beta(1,1)          8     0.7778                   1.286\n",
+      "uniform Beta(1,1)         16     0.8824                   1.133\n",
+      "uniform Beta(1,1)         32     0.9394                   1.065\n",
+      "easy-heavy Beta(2,0.5)     8     0.5651                   1.770\n",
+      "easy-heavy Beta(2,0.5)    16     0.6819                   1.467\n",
+      "easy-heavy Beta(2,0.5)    32     0.7705                   1.298\n",
+      "late RL Beta(4,0.5)        8     0.4338                   2.305\n",
+      "late RL Beta(4,0.5)       16     0.5637                   1.774\n",
+      "late RL Beta(4,0.5)       32     0.6758                   1.480\n"
+     ]
+    }
+   ],
+   "source": [
+    "def f_info(a, b, G):\n",
+    "    return 1 - math.exp(betaln(a + G, b) - betaln(a, b)) - math.exp(betaln(a, b + G) - betaln(a, b))\n",
+    "rs = np.random.RandomState(9)\n",
+    "worst = 0\n",
+    "for (a, b, G_) in [(0.5, 1.0, 8), (2.0, 0.5, 16), (1.0, 1.0, 8)]:\n",
+    "    p = rs.beta(a, b, size=4_000_000); fmc = 1 - (p ** G_ + (1 - p) ** G_).mean()\n",
+    "    se = (p ** G_ + (1 - p) ** G_).std() / math.sqrt(len(p))\n",
+    "    worst = max(worst, abs(fmc - f_info(a, b, G_)) / se)\n",
+    "check('7. informative fraction: Beta-function closed form vs MC (max z)', worst, 4.0)\n",
+    "print('population            G    informative fraction f   oversampling factor 1/f')\n",
+    "for (nm, a, b) in [('hard-heavy Beta(0.5,1)', 0.5, 1.0), ('uniform Beta(1,1)', 1.0, 1.0), ('easy-heavy Beta(2,0.5)', 2.0, 0.5), ('late RL Beta(4,0.5)', 4.0, 0.5)]:\n",
+    "    for G_ in (8, 16, 32):\n",
+    "        f = f_info(a, b, G_); print(f'{nm:24s} {G_:3d}   {f:8.4f}                 {1/f:7.3f}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d7e59883",
+   "metadata": {},
+   "source": [
+    "## Verification table"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "0be2bb80",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:50:21.137898Z",
+     "iopub.status.busy": "2026-10-14T09:50:21.137739Z",
+     "iopub.status.idle": "2026-10-14T09:50:21.141267Z",
+     "shell.execute_reply": "2026-10-14T09:50:21.140434Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "identity                                                                                      max abs err\n",
+      "1. autograd surrogate gradient = p_hat (1 - p_hat), all c in 0..8, four p                       2.776e-17\n",
+      "2. E[p_hat(1-p_hat)] = (1-1/G) p(1-p), exact rationals                                          0.000e+00\n",
+      "2. max of p(1-p) is 1/4 at p = 1/2                                                              0.000e+00\n",
+      "3. MC mean of p_hat(1-p_hat) vs exact (in s.e.)                                                 7.820e-01\n",
+      "3. MC dead-group fraction vs p^G + (1-p)^G (in s.e.)                                            3.515e-01\n",
+      "4. std-normalised group gradient = sigma_hat (direct sum vs closed form)                        1.110e-16\n",
+      "4. unbiased-std normalisation: group gradient = sigma_hat sqrt((G-1)/G)                         1.110e-16\n",
+      "4. std-normalisation re-weights c=1 vs c=4 groups by exactly sigma(4)/sigma(1)                  2.220e-16\n",
+      "5. Cauchy-Schwarz bound ||grad p||^2 <= p(1-p) E||s||^2 (max violation of ratio <= 1)           0.000e+00\n",
+      "6. band weight w_tau(p): closed form vs simulation (p=.7,G=8,tau=.75)                           3.628e-04\n",
+      "6. no band (tau=1) gives exactly 1-1/G for all p                                                3.331e-16\n",
+      "7. informative fraction: Beta-function closed form vs MC (max z)                                9.161e-01\n",
+      "CHECKS_JSON:[[\"1. autograd surrogate gradient = p_hat (1 - p_hat), all c in 0..8, four p\", 2.7755575615628914e-17, 1e-15], [\"2. E[p_hat(1-p_hat)] = (1-1/G) p(1-p), exact rationals\", 0.0, 0.0], [\"2. max of p(1-p) is 1/4 at p = 1/2\", 0.0, 1e-12], [\"3. MC mean of p_hat(1-p_hat) vs exact (in s.e.)\", 0.7820444294887922, 4.0], [\"3. MC dead-group fraction vs p^G + (1-p)^G (in s.e.)\", 0.3515169230631549, 4.0], [\"4. std-normalised group gradient = sigma_hat (direct sum vs closed form)\", 1.1102230246251565e-16, 1e-15], [\"4. unbiased-std normalisation: group gradient = sigma_hat sqrt((G-1)/G)\", 1.1102230246251565e-16, 1e-15], [\"4. std-normalisation re-weights c=1 vs c=4 groups by exactly sigma(4)/sigma(1)\", 2.220446049250313e-16, 1e-14], [\"5. Cauchy-Schwarz bound ||grad p||^2 <= p(1-p) E||s||^2 (max violation of ratio <= 1)\", 0.0, 0.0], [\"6. band weight w_tau(p): closed form vs simulation (p=.7,G=8,tau=.75)\", 0.0003628045634921362, 0.003], [\"6. no band (tau=1) gives exactly 1-1/G for all p\", 3.3306690738754696e-16, 1e-14], [\"7. informative fraction: Beta-function closed form vs MC (max z)\", 0.9161417930924225, 4.0]]\n"
+     ]
+    }
+   ],
+   "source": [
+    "print()\n",
+    "print(f'{\"identity\":92s} {\"max abs err\":>12s}')\n",
+    "for n_, e_, t_ in CHECKS: print(f'{n_:92s} {e_:12.3e}')\n",
+    "print('CHECKS_JSON:' + json.dumps(CHECKS))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

code/notebooks/0207-test-time-scaling-cost-model.ipynb

@@ -0,0 +1,602 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "b8fba5aa",
+   "metadata": {},
+   "source": [
+    "# 0207 -- Test-time scaling: pass@k vs compute, self-correction rounds vs more samples, best-of-n with a perfect verifier\n",
+    "\n",
+    "Identities verified (exact rationals / fp64):\n",
+    "1. Best-of-n with a perfect verifier: `coverage(n) = 1 - E[(1-p)^n]`; with `p ~ Beta(a,b)`: `1 - B(a,b+n)/B(a,b)`, and `1 - coverage ~ Gamma(a+b)/Gamma(b) n^-a`.\n",
+    "2. Early stopping: `E[min(N,n)] = (1-(1-p)^n)/p`; Beta-average `sum_{j<n} B(a,b+j)/B(a,b)`.\n",
+    "3. Posterior hazard `P(success on attempt j+1 | j failures) = a/(a+b+j)`; optimal stopping by dynamic programming = threshold rule `j* = a U/c - (a+b)`.\n",
+    "4. Equal-cost comparison: adaptive stopping vs fixed n.\n",
+    "5. Fresh samples vs revisions: revise iff `ln(1-q)/ln(1-p_1) >= rho` (brute-force enumeration over `(n,R)`).\n",
+    "6. Imperfect verifier: precision cap `p/(p + (1-p) phi)` of the first accepted sample, independent of n.\n",
+    "7. The exponentiated power law `c = exp(a k^b)` (Large Language Monkeys) vs the Beta-mixture coverage curve."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "d1a9b6b8",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:51:01.379846Z",
+     "iopub.status.busy": "2026-10-14T09:51:01.379697Z",
+     "iopub.status.idle": "2026-10-14T09:51:01.728451Z",
+     "shell.execute_reply": "2026-10-14T09:51:01.727516Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n",
+    "import math, json, itertools\n",
+    "from fractions import Fraction as Fr\n",
+    "from math import comb, lgamma\n",
+    "import numpy as np\n",
+    "from scipy.special import betaln, gammaln\n",
+    "from scipy.optimize import least_squares\n",
+    "CHECKS = []\n",
+    "def check(name, err, tol):\n",
+    "    CHECKS.append((name, float(err), float(tol)))\n",
+    "    assert err <= tol, (name, float(err), tol)\n",
+    "    print(f'{name:92s} err={float(err):.3e}  (tol {tol:.0e})')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f09fd568",
+   "metadata": {},
+   "source": [
+    "## 1. Coverage with a perfect verifier\n",
+    "NAIVE: Monte Carlo over problems (draw `p`, draw `n` attempts, any success?). EFFICIENT: Beta-function closed form."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "57c4d2e9",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:51:01.730089Z",
+     "iopub.status.busy": "2026-10-14T09:51:01.729862Z",
+     "iopub.status.idle": "2026-10-14T09:51:03.100278Z",
+     "shell.execute_reply": "2026-10-14T09:51:03.099451Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1. coverage: Monte Carlo vs 1 - B(a,b+n)/B(a,b)  (max z, n in {1..256})                      err=1.274e+00  (tol 4e+00)\n",
+      "(1-coverage) / [Gamma(a+b)/Gamma(b) n^-a] at n = [10, 100, 1000, 10000, 100000] -> [0.939, 0.9933, 0.9993, 0.9999, 1.0]\n",
+      "1. tail law: ratio -> 1 as n -> 1e5 (|ratio-1|)                                              err=6.800e-06  (tol 5e-03)\n"
+     ]
+    }
+   ],
+   "source": [
+    "def coverage_beta(a, b, n): return 1 - np.exp(betaln(a, b + n) - betaln(a, b))\n",
+    "rs = np.random.RandomState(0)\n",
+    "a, b = 0.4, 2.0\n",
+    "worst = 0.0\n",
+    "for n in (1, 4, 16, 64, 256):\n",
+    "    p = rs.beta(a, b, size=2_000_000); hit = rs.rand(len(p)) < 1 - (1 - p) ** n\n",
+    "    se = math.sqrt(hit.mean() * (1 - hit.mean()) / len(p))\n",
+    "    worst = max(worst, abs(hit.mean() - coverage_beta(a, b, n)) / se)\n",
+    "check('1. coverage: Monte Carlo vs 1 - B(a,b+n)/B(a,b)  (max z, n in {1..256})', worst, 4.0)\n",
+    "ns = np.array([10, 100, 1000, 10_000, 100_000])\n",
+    "ratio = (1 - coverage_beta(a, b, ns)) / (np.exp(gammaln(a + b) - gammaln(b)) * ns ** (-a))\n",
+    "print('(1-coverage) / [Gamma(a+b)/Gamma(b) n^-a] at n =', ns.tolist(), '->', np.round(ratio, 4).tolist())\n",
+    "check('1. tail law: ratio -> 1 as n -> 1e5 (|ratio-1|)', abs(ratio[-1] - 1), 5e-3)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b908b5d2",
+   "metadata": {},
+   "source": [
+    "## 2. Early stopping with a perfect verifier\n",
+    "Sample sequentially, stop at the first verified success or after `n` attempts. Cost `N_n = min(Geom(p), n)`, `E[N_n] = sum_{j=0}^{n-1} (1-p)^j = (1-(1-p)^n)/p`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "f4f195cb",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:51:03.101782Z",
+     "iopub.status.busy": "2026-10-14T09:51:03.101629Z",
+     "iopub.status.idle": "2026-10-14T09:51:03.424265Z",
+     "shell.execute_reply": "2026-10-14T09:51:03.423468Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2. E[min(Geom(p),n)] = (1-(1-p)^n)/p = sum_{j<n}(1-p)^j  (exact rationals)                   err=0.000e+00  (tol 0e+00)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2. Beta-average of the early-stopping cost: MC vs sum_j B(a,b+j)/B(a,b)  (z)                 err=4.881e-01  (tol 4e+00)\n",
+      "Beta(0.4,2.0), n=64: coverage=0.7671; expected attempts per problem with early stopping = 23.05 (vs 64 without)\n"
+     ]
+    }
+   ],
+   "source": [
+    "def mean_cost_exact(p, n):\n",
+    "    p = Fr(p); return sum((1 - p) ** j for j in range(n))\n",
+    "worst = 0.0\n",
+    "for p in (Fr(1, 100), Fr(1, 7), Fr(1, 2)):\n",
+    "    for n in (1, 5, 40):\n",
+    "        # NAIVE: enumerate the distribution of min(Geom, n) explicitly\n",
+    "        e = sum(Fr(k) * (1 - p) ** (k - 1) * p for k in range(1, n)) + Fr(n) * (1 - p) ** (n - 1)\n",
+    "        worst = max(worst, float(abs(e - (1 - (1 - p) ** n) / p)), float(abs(e - mean_cost_exact(p, n))))\n",
+    "check('2. E[min(Geom(p),n)] = (1-(1-p)^n)/p = sum_{j<n}(1-p)^j  (exact rationals)', worst, 0.0)\n",
+    "def cost_beta(a, b, n): return sum(math.exp(betaln(a, b + j) - betaln(a, b)) for j in range(n))\n",
+    "rs = np.random.RandomState(1); a, b = 0.4, 2.0; n = 64\n",
+    "p = rs.beta(a, b, size=2_000_000); N_ = np.minimum(rs.geometric(np.clip(p, 1e-300, 1)), n)\n",
+    "check('2. Beta-average of the early-stopping cost: MC vs sum_j B(a,b+j)/B(a,b)  (z)', abs(N_.mean() - cost_beta(a, b, n)) / (N_.std() / math.sqrt(len(N_))), 4.0)\n",
+    "print(f'Beta({a},{b}), n={n}: coverage={coverage_beta(a,b,n):.4f}; expected attempts per problem with early stopping = {cost_beta(a,b,n):.2f} (vs {n} without)')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "489971ea",
+   "metadata": {},
+   "source": [
+    "## 3. Posterior hazard and the optimal stopping rule\n",
+    "For `p ~ Beta(a,b)` the posterior after `j` failures is `Beta(a, b+j)`, so `h_j = P(success | j failures) = a/(a+b+j)`.\n",
+    "Each attempt costs `c`; a solved problem is worth `U`. DP: `W_j = max(0, h_j (U - c) - (1-h_j) c + (1-h_j) W_{j+1})` with `W_{n_max}=0`.\n",
+    "Because `h_j` is decreasing the one-step-look-ahead rule is optimal: continue iff `h_j U >= c`, i.e. `j < a U/c - (a+b)`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "087dfa00",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:51:03.425743Z",
+     "iopub.status.busy": "2026-10-14T09:51:03.425591Z",
+     "iopub.status.idle": "2026-10-14T09:51:05.083274Z",
+     "shell.execute_reply": "2026-10-14T09:51:05.082473Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "3. posterior success probability after j failures = a/(a+b+j)  (importance-sampling MC, max z) err=1.647e+00  (tol 4e+00)\n",
+      "3. same, exact Beta-function ratio B(a+1,b+j)/B(a,b+j)                                       err=8.491e-16  (tol 1e-14)\n",
+      "U=10.0, c=0.05: DP stops after 78 attempts; threshold a U/c - (a+b) = 77.60 -> ceil 78;  value W_0=6.5345\n",
+      "U=10.0, c=0.01: DP stops after 398 attempts; threshold a U/c - (a+b) = 397.60 -> ceil 398;  value W_0=8.1394\n",
+      "U=50.0, c=0.05: DP stops after 398 attempts; threshold a U/c - (a+b) = 397.60 -> ceil 398;  value W_0=40.6968\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "3. DP stopping index vs threshold rule ceil(aU/c-(a+b)) (max |diff| over 3 settings)         err=0.000e+00  (tol 1e+00)\n"
+     ]
+    }
+   ],
+   "source": [
+    "a, b = 0.4, 2.0\n",
+    "# hazard check by self-normalised importance sampling from the prior (weights (1-p)^j = likelihood of j failures)\n",
+    "rs = np.random.RandomState(7); pp = rs.beta(a, b, size=8_000_000); worst_z = 0.0\n",
+    "for j in (0, 1, 5, 20):\n",
+    "    w = (1 - pp) ** j; est = (pp * w).sum() / w.sum()\n",
+    "    # delta-method s.e. of a ratio estimator\n",
+    "    se = math.sqrt((((pp - est) * w) ** 2).sum()) / w.sum()\n",
+    "    worst_z = max(worst_z, abs(est - a / (a + b + j)) / se)\n",
+    "check('3. posterior success probability after j failures = a/(a+b+j)  (importance-sampling MC, max z)', worst_z, 4.5)\n",
+    "# exact version via Beta functions (no quadrature):\n",
+    "worst = max(abs(math.exp(betaln(a + 1, b + j) - betaln(a, b + j)) - a / (a + b + j)) for j in range(0, 200))\n",
+    "check('3. same, exact Beta-function ratio B(a+1,b+j)/B(a,b+j)', worst, 1e-14)\n",
+    "def dp_value(U, c, a, b, nmax=100000):\n",
+    "    W = 0.0; cont = []\n",
+    "    for j in range(nmax - 1, -1, -1):\n",
+    "        h = a / (a + b + j)\n",
+    "        val = h * (U - c) - (1 - h) * c + (1 - h) * W\n",
+    "        W = max(0.0, val); cont.append(val > 0)\n",
+    "    cont = cont[::-1]\n",
+    "    jstop = next((j for j, v in enumerate(cont) if not v), nmax)\n",
+    "    return W, jstop\n",
+    "for (U, c) in [(10.0, 0.05), (10.0, 0.01), (50.0, 0.05)]:\n",
+    "    W, jstop = dp_value(U, c, a, b)\n",
+    "    j_theory = math.ceil(a * U / c - (a + b))\n",
+    "    print(f'U={U}, c={c}: DP stops after {jstop} attempts; threshold a U/c - (a+b) = {a*U/c-(a+b):.2f} -> ceil {j_theory};  value W_0={W:.4f}')\n",
+    "    assert abs(jstop - j_theory) <= 1\n",
+    "check('3. DP stopping index vs threshold rule ceil(aU/c-(a+b)) (max |diff| over 3 settings)', max(abs(dp_value(U, c, a, b)[1] - math.ceil(a * U / c - (a + b))) for U, c in [(10.0, 0.05), (10.0, 0.01), (50.0, 0.05)]), 1)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6e820e32",
+   "metadata": {},
+   "source": [
+    "## 4. Equal expected cost: adaptive stopping vs a fixed budget\n",
+    "Fixed `n = 64` attempts for every problem costs 64 each. With early stopping the expected cost is `cost(n_max)`; choose `n_max` so that it equals 64 and compare coverage."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "bddd003b",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:51:05.085063Z",
+     "iopub.status.busy": "2026-10-14T09:51:05.084733Z",
+     "iopub.status.idle": "2026-10-14T09:51:19.832181Z",
+     "shell.execute_reply": "2026-10-14T09:51:19.831457Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Beta(0.4,2.0): fixed n=64 -> coverage 0.7671;  adaptive with the same mean cost 64 needs n_max = 323 -> coverage 0.8771\n",
+      "4. n_max is the smallest n with expected cost >= 64 (violations; 0 = cost(n_max-1) < 64 <= cost(n_max)) err=0.000e+00  (tol 0e+00)\n",
+      "cost at n_max-1, n_max: 63.997038900831 64.1201008159222\n"
+     ]
+    }
+   ],
+   "source": [
+    "target = 64.0\n",
+    "# cumulative expected cost cost(n) = sum_{j<n} B(a,b+j)/B(a,b) via a vectorised cumulative sum\n",
+    "j = np.arange(0, 20_000_000); terms = np.exp(betaln(a, b + j) - betaln(a, b)); cs = np.cumsum(terms)\n",
+    "n_adapt = int(np.searchsorted(cs, target)) + 1\n",
+    "print(f'Beta({a},{b}): fixed n=64 -> coverage {coverage_beta(a,b,64):.4f};  adaptive with the same mean cost 64 needs n_max = {n_adapt} -> coverage {coverage_beta(a,b,n_adapt):.4f}')\n",
+    "check('4. n_max is the smallest n with expected cost >= 64 (violations; 0 = cost(n_max-1) < 64 <= cost(n_max))', float(not (cs[n_adapt - 2] < target <= cs[n_adapt - 1])), 0.0)\n",
+    "print('cost at n_max-1, n_max:', cs[n_adapt - 2], cs[n_adapt - 1])\n",
+    "assert coverage_beta(a, b, n_adapt) > coverage_beta(a, b, 64)\n",
+    "CHECKS.append(('4. adaptive-minus-fixed coverage at equal mean cost (informational)', float(coverage_beta(a, b, n_adapt) - coverage_beta(a, b, 64)), float('inf')))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "867deb08",
+   "metadata": {},
+   "source": [
+    "## 5. Fresh samples vs revisions at equal token cost\n",
+    "A fresh sample costs 1 unit and succeeds w.p. `p1`. A chain is one attempt followed by up to `R` revisions, each costing `rho` units and succeeding w.p. `q` given all earlier failures\n",
+    "(independent of the earlier attempts: the simplest model). Chain failure `(1-p1)(1-q)^R`, cost `1 + rho R`. With budget `B` we run `n` independent chains.\n",
+    "Failure `= [(1-p1)(1-q)^R]^n`. Log-failure per cost: `f(R) = (a1 + R a2)/(1+rho R)`, `a1 = -ln(1-p1)`, `a2 = -ln(1-q)`. `f' ∝ a2 - rho a1`: revise iff `a2 > rho a1`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "d391bbf9",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:51:19.833654Z",
+     "iopub.status.busy": "2026-10-14T09:51:19.833497Z",
+     "iopub.status.idle": "2026-10-14T09:51:19.840774Z",
+     "shell.execute_reply": "2026-10-14T09:51:19.840044Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(p1, q, rho, revise predicted, best n, best R, failure prob. of best mix at budget 600, failure prob. with fresh samples only)\n",
+      "  (0.05, 0.03, 1.5, False, 600, 0, 4.31e-14, 4.31e-14)\n",
+      "  (0.05, 0.06, 1.0, True, 50, 11, 1.28e-16, 4.31e-14)\n",
+      "  (0.05, 0.06, 1.5, False, 600, 0, 4.31e-14, 4.31e-14)\n",
+      "  (0.1, 0.12, 1.0, True, 50, 11, 1.51e-33, 3.51e-28)\n",
+      "sign rule vs brute-force optimum: agree on 36/36 grid points\n",
+      "disagreements (p1, q, rho, relative gap |a2 - rho a1| / (rho a1), best (n,R)): []\n",
+      "5. revise-iff-(a2 > rho a1) rule vs brute-force optimum: disagreements with relative gap > 10% err=0.000e+00  (tol 0e+00)\n"
+     ]
+    }
+   ],
+   "source": [
+    "def best_mix(p1, q, rho, B, Rmax=12):\n",
+    "    \"\"\"Minimise the log failure probability n*[ln(1-p1) + R ln(1-q)] over integer (n, R) with n(1+rho R) <= B (avoids float saturation at success ~ 1).\"\"\"\n",
+    "    best = (float('inf'), None)\n",
+    "    for R in range(0, Rmax + 1):\n",
+    "        n = int(B // (1 + rho * R))\n",
+    "        if n < 1: continue\n",
+    "        logfail = n * (math.log(1 - p1) + R * math.log(1 - q))\n",
+    "        if logfail < best[0] - 1e-12: best = (logfail, (n, R))\n",
+    "    return math.exp(best[0]), best[1]        # failure probability of the best mix\n",
+    "rows = []; borderline = []\n",
+    "agree = 0; total = 0\n",
+    "for p1 in (0.02, 0.05, 0.1):\n",
+    "    for q in (0.01, 0.03, 0.06, 0.12):\n",
+    "        for rho in (1.0, 1.5, 2.5):\n",
+    "            B = 600.0\n",
+    "            succ, (n_, R_) = best_mix(p1, q, rho, B)   # `succ` is the failure probability here\n",
+    "            predicted_revise = (-math.log(1 - q)) > rho * (-math.log(1 - p1))\n",
+    "            gap = abs((-math.log(1 - q)) - rho * (-math.log(1 - p1))) / (rho * (-math.log(1 - p1)))\n",
+    "            total += 1\n",
+    "            if (R_ > 0) == predicted_revise: agree += 1\n",
+    "            else:\n",
+    "                borderline.append((p1, q, rho, round(gap, 3), (n_, R_)))\n",
+    "            if (p1, q, rho) in [(0.05, 0.06, 1.0), (0.05, 0.06, 1.5), (0.05, 0.03, 1.5), (0.1, 0.12, 1.0)]:\n",
+    "                rows.append((p1, q, rho, predicted_revise, n_, R_, float('%.3g' % succ), float('%.3g' % ((1 - p1) ** int(B)))))\n",
+    "print('(p1, q, rho, revise predicted, best n, best R, failure prob. of best mix at budget 600, failure prob. with fresh samples only)')\n",
+    "for r in rows: print(' ', r)\n",
+    "print(f'sign rule vs brute-force optimum: agree on {agree}/{total} grid points')\n",
+    "print('disagreements (p1, q, rho, relative gap |a2 - rho a1| / (rho a1), best (n,R)):', borderline)\n",
+    "far = [b_ for b_ in borderline if b_[3] > 0.1]\n",
+    "check('5. revise-iff-(a2 > rho a1) rule vs brute-force optimum: disagreements with relative gap > 10%', len(far), 0)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b061849f",
+   "metadata": {},
+   "source": [
+    "## 6. Imperfect verifier: the first *accepted* sample is correct only with probability `p/(p+(1-p) phi)`\n",
+    "Verifier accepts correct proofs always and incorrect ones w.p. `phi` (false accept). Sampling until the first accept, the accepted sample is correct w.p. `p/(p + (1-p) phi)`,\n",
+    "for any budget `n` (precision cap), while a perfect verifier gives 1 with the same stopping rule."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "d216a73d",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:51:19.842276Z",
+     "iopub.status.busy": "2026-10-14T09:51:19.842130Z",
+     "iopub.status.idle": "2026-10-14T09:51:24.279548Z",
+     "shell.execute_reply": "2026-10-14T09:51:24.278840Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "p=0.05, false accept=0.01: precision of an accepted sample = 0.8414  (theory 0.8403)\n",
+      "p=0.2, false accept=0.05: precision of an accepted sample = 0.8335  (theory 0.8333)\n",
+      "p=0.01, false accept=0.002: precision of an accepted sample = 0.8315  (theory 0.8347)\n",
+      "6. precision of the first accepted sample: MC vs p/(p+(1-p)phi)  (max z)                     err=1.883e+00  (tol 4e+00)\n"
+     ]
+    }
+   ],
+   "source": [
+    "rs = np.random.RandomState(4)\n",
+    "worst = 0\n",
+    "for (p, phi) in [(0.05, 0.01), (0.2, 0.05), (0.01, 0.002)]:\n",
+    "    N = 4_000_000\n",
+    "    corr = rs.rand(N) < p; acc = corr | (rs.rand(N) < phi)\n",
+    "    precision = corr[acc].mean(); se = math.sqrt(precision * (1 - precision) / acc.sum())\n",
+    "    worst = max(worst, abs(precision - p / (p + (1 - p) * phi)) / se)\n",
+    "    print(f'p={p}, false accept={phi}: precision of an accepted sample = {precision:.4f}  (theory {p/(p+(1-p)*phi):.4f})')\n",
+    "check('6. precision of the first accepted sample: MC vs p/(p+(1-p)phi)  (max z)', worst, 4.0)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7215e58c",
+   "metadata": {},
+   "source": [
+    "## 7. The empirical exponentiated power law vs the Beta-mixture curve\n",
+    "Large Language Monkeys fit coverage `c(k) = exp(a k^b)` (Fig. 5). For `p ~ Beta(alpha, beta)` the exact curve is `1 - B(alpha, beta+k)/B(alpha,beta)`. Fit `log c` by least squares on log-spaced `k` in `[1, 10^4]`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "32ba1c14",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:51:24.280954Z",
+     "iopub.status.busy": "2026-10-14T09:51:24.280801Z",
+     "iopub.status.idle": "2026-10-14T09:51:24.302130Z",
+     "shell.execute_reply": "2026-10-14T09:51:24.301470Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Beta(0.3,1.5): fitted c=exp(a k^b): a=-1.648, b=-0.361, mean abs error=0.0061, max=0.0259\n",
+      "Beta(0.5,2.0): fitted c=exp(a k^b): a=-1.673, b=-0.535, mean abs error=0.0014, max=0.0123\n",
+      "Beta(1.0,3.0): fitted c=exp(a k^b): a=-1.625, b=-0.809, mean abs error=0.0062, max=0.0530\n"
+     ]
+    }
+   ],
+   "source": [
+    "ks = np.unique(np.logspace(0, 4, 100).astype(int))\n",
+    "for (al, be) in [(0.3, 1.5), (0.5, 2.0), (1.0, 3.0)]:\n",
+    "    cov = coverage_beta(al, be, ks)\n",
+    "    fit = least_squares(lambda th: np.exp(th[0] * ks ** th[1]) - cov, x0=[-1.0, -0.3], bounds=([-20, -3], [-1e-6, -1e-6]))\n",
+    "    print(f'Beta({al},{be}): fitted c=exp(a k^b): a={fit.x[0]:.3f}, b={fit.x[1]:.3f}, mean abs error={np.abs(fit.fun).mean():.4f}, max={np.abs(fit.fun).max():.4f}')\n",
+    "CHECKS.append(('7. exp-power-law fit error on a Beta-mixture curve (max abs over 3 fits, informational)', float(max(np.abs(least_squares(lambda th, cv=coverage_beta(al, be, ks): np.exp(th[0] * ks ** th[1]) - cv, x0=[-1.0, -0.3], bounds=([-20, -3], [-1e-6, -1e-6])).fun).max() for al, be in [(0.3, 1.5), (0.5, 2.0), (1.0, 3.0)])), float('inf')))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f9f47176",
+   "metadata": {},
+   "source": [
+    "## 8. Token accounting for a published comparison (arithmetic only)\n",
+    "Goedel-Prover-V2 Sec. 3.2: first round capped at 30,000 tokens, self-correction mode (initial + 2 revisions) at 40,000 tokens in total.\n",
+    "Sec. 3.5: with the context extended to 128k (YaRN) and up to 5 revisions, pass@32 with correction reaches 92.7% > 92.2% (32B, pass@8192, no correction).\n",
+    "Table 4: the 40k-token self-correction mode reaches only 90.4% at pass@32. Ratio of caps (neither an upper nor a lower bound on tokens actually used):"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "cbf4ed2e",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:51:24.303439Z",
+     "iopub.status.busy": "2026-10-14T09:51:24.303292Z",
+     "iopub.status.idle": "2026-10-14T09:51:24.307308Z",
+     "shell.execute_reply": "2026-10-14T09:51:24.306689Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "cap ratio, 40k self-correction mode (reaches 90.4% < 92.2%): 8192 x 30k / (32 x 40k)  = 192x\n",
+      "cap ratio, 128k extended mode (reaches 92.7% > 92.2%)      : 8192 x 30k / (32 x 128k) = 60x\n",
+      "rho implied by the caps (chain = attempt + 2 revisions <= 40k vs fresh <= 30k): 0.16666666666666663\n",
+      "8. rho implied by 1 + 2 rho = 40k/30k equals 1/6                                             err=2.776e-17  (tol 1e-15)\n"
+     ]
+    }
+   ],
+   "source": [
+    "tok_std, tok_sc40, tok_sc128 = 30_000, 40_000, 128_000\n",
+    "ratio40 = (8192 * tok_std) / (32 * tok_sc40); ratio128 = (8192 * tok_std) / (32 * tok_sc128)\n",
+    "print(f'cap ratio, 40k self-correction mode (reaches 90.4% < 92.2%): 8192 x 30k / (32 x 40k)  = {ratio40:.0f}x')\n",
+    "print(f'cap ratio, 128k extended mode (reaches 92.7% > 92.2%)      : 8192 x 30k / (32 x 128k) = {ratio128:.0f}x')\n",
+    "assert ratio40 == 192 and ratio128 == 60\n",
+    "# chain-vs-fresh cost at the caps: 1 + 2 rho <= 40/30  =>  rho <= 1/6 if the caps were costs\n",
+    "rho_cap = (tok_sc40 / tok_std - 1) / 2\n",
+    "print('rho implied by the caps (chain = attempt + 2 revisions <= 40k vs fresh <= 30k):', rho_cap)\n",
+    "check('8. rho implied by 1 + 2 rho = 40k/30k equals 1/6', abs(rho_cap - 1 / 6), 1e-15)\n",
+    "CHECKS.append(('8. cap-based token ratio for the 92.7% (128k, <=5 revisions) configuration (informational)', float(ratio128), float('inf')))\n",
+    "CHECKS.append(('8. cap-based token ratio for the 40k self-correction mode (informational; that mode does not beat 92.2%)', float(ratio40), float('inf')))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d23da596",
+   "metadata": {},
+   "source": [
+    "## 8b. Sensitivity of the brute-force optimum to the chain-length cap\n",
+    "With `q` independent of history the continuous objective is monotone in `R`; the integer optimum depends on the cap `R_max` (and on `floor(B/(1+rho R))`)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "56d4dfdb",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:51:24.308569Z",
+     "iopub.status.busy": "2026-10-14T09:51:24.308427Z",
+     "iopub.status.idle": "2026-10-14T09:51:24.313929Z",
+     "shell.execute_reply": "2026-10-14T09:51:24.313186Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "p1=0.05 q=0.06 rho=1 B=600, R_max= 12: best (n,R)=(50,11), log-failure = -36.596\n",
+      "p1=0.05 q=0.06 rho=1 B=600, R_max= 50: best (n,R)=(12,49), log-failure = -36.998\n",
+      "p1=0.05 q=0.06 rho=1 B=600, R_max=599: best (n,R)=(1,599), log-failure = -37.115\n"
+     ]
+    }
+   ],
+   "source": [
+    "def best_mix_cap(p1, q, rho, B, Rmax):\n",
+    "    best = (float('inf'), None)\n",
+    "    for R in range(0, Rmax + 1):\n",
+    "        n = int(B // (1 + rho * R))\n",
+    "        if n < 1: continue\n",
+    "        lf = n * (math.log(1 - p1) + R * math.log(1 - q))\n",
+    "        if lf < best[0] - 1e-12: best = (lf, (n, R))\n",
+    "    return best\n",
+    "for cap in (12, 50, 599):\n",
+    "    lf, (n_, R_) = best_mix_cap(0.05, 0.06, 1.0, 600.0, cap)\n",
+    "    print(f'p1=0.05 q=0.06 rho=1 B=600, R_max={cap:3d}: best (n,R)=({n_},{R_}), log-failure = {lf:.3f}')\n",
+    "lf12 = best_mix_cap(0.05, 0.06, 1.0, 600.0, 12); lf599 = best_mix_cap(0.05, 0.06, 1.0, 600.0, 599)\n",
+    "assert lf12[1] == (50, 11) and lf599[1][0] <= 2 and lf599[0] < lf12[0]\n",
+    "CHECKS.append(('8b. optimum for R_max=12 is (50,11); uncapped optimum has n<=2 chains (informational)', float(lf599[1][0]), float('inf')))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b5ac6a6d",
+   "metadata": {},
+   "source": [
+    "## Verification table"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "433dc5ca",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:51:24.315222Z",
+     "iopub.status.busy": "2026-10-14T09:51:24.315080Z",
+     "iopub.status.idle": "2026-10-14T09:51:24.318768Z",
+     "shell.execute_reply": "2026-10-14T09:51:24.317925Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "identity                                                                                        max abs err\n",
+      "1. coverage: Monte Carlo vs 1 - B(a,b+n)/B(a,b)  (max z, n in {1..256})                           1.274e+00\n",
+      "1. tail law: ratio -> 1 as n -> 1e5 (|ratio-1|)                                                   6.800e-06\n",
+      "2. E[min(Geom(p),n)] = (1-(1-p)^n)/p = sum_{j<n}(1-p)^j  (exact rationals)                        0.000e+00\n",
+      "2. Beta-average of the early-stopping cost: MC vs sum_j B(a,b+j)/B(a,b)  (z)                      4.881e-01\n",
+      "3. posterior success probability after j failures = a/(a+b+j)  (importance-sampling MC, max z)    1.647e+00\n",
+      "3. same, exact Beta-function ratio B(a+1,b+j)/B(a,b+j)                                            8.491e-16\n",
+      "3. DP stopping index vs threshold rule ceil(aU/c-(a+b)) (max |diff| over 3 settings)              0.000e+00\n",
+      "4. n_max is the smallest n with expected cost >= 64 (violations; 0 = cost(n_max-1) < 64 <= cost(n_max))    0.000e+00\n",
+      "4. adaptive-minus-fixed coverage at equal mean cost (informational)                               1.100e-01\n",
+      "5. revise-iff-(a2 > rho a1) rule vs brute-force optimum: disagreements with relative gap > 10%    0.000e+00\n",
+      "6. precision of the first accepted sample: MC vs p/(p+(1-p)phi)  (max z)                          1.883e+00\n",
+      "7. exp-power-law fit error on a Beta-mixture curve (max abs over 3 fits, informational)           5.301e-02\n",
+      "8. rho implied by 1 + 2 rho = 40k/30k equals 1/6                                                  2.776e-17\n",
+      "8. cap-based token ratio for the 92.7% (128k, <=5 revisions) configuration (informational)        6.000e+01\n",
+      "8. cap-based token ratio for the 40k self-correction mode (informational; that mode does not beat 92.2%)    1.920e+02\n",
+      "8b. optimum for R_max=12 is (50,11); uncapped optimum has n<=2 chains (informational)             1.000e+00\n",
+      "CHECKS_JSON:[[\"1. coverage: Monte Carlo vs 1 - B(a,b+n)/B(a,b)  (max z, n in {1..256})\", 1.2744892726318469, 4.0], [\"1. tail law: ratio -> 1 as n -> 1e5 (|ratio-1|)\", 6.8000623573150065e-06, 0.005], [\"2. E[min(Geom(p),n)] = (1-(1-p)^n)/p = sum_{j<n}(1-p)^j  (exact rationals)\", 0.0, 0.0], [\"2. Beta-average of the early-stopping cost: MC vs sum_j B(a,b+j)/B(a,b)  (z)\", 0.488083889788271, 4.0], [\"3. posterior success probability after j failures = a/(a+b+j)  (importance-sampling MC, max z)\", 1.6465685132559886, 4.5], [\"3. same, exact Beta-function ratio B(a+1,b+j)/B(a,b+j)\", 8.491471414906471e-16, 1e-14], [\"3. DP stopping index vs threshold rule ceil(aU/c-(a+b)) (max |diff| over 3 settings)\", 0.0, 1.0], [\"4. n_max is the smallest n with expected cost >= 64 (violations; 0 = cost(n_max-1) < 64 <= cost(n_max))\", 0.0, 0.0], [\"4. adaptive-minus-fixed coverage at equal mean cost (informational)\", 0.10998111815038014, Infinity], [\"5. revise-iff-(a2 > rho a1) rule vs brute-force optimum: disagreements with relative gap > 10%\", 0.0, 0.0], [\"6. precision of the first accepted sample: MC vs p/(p+(1-p)phi)  (max z)\", 1.88315480662085, 4.0], [\"7. exp-power-law fit error on a Beta-mixture curve (max abs over 3 fits, informational)\", 0.05301359850390941, Infinity], [\"8. rho implied by 1 + 2 rho = 40k/30k equals 1/6\", 2.7755575615628914e-17, 1e-15], [\"8. cap-based token ratio for the 92.7% (128k, <=5 revisions) configuration (informational)\", 60.0, Infinity], [\"8. cap-based token ratio for the 40k self-correction mode (informational; that mode does not beat 92.2%)\", 192.0, Infinity], [\"8b. optimum for R_max=12 is (50,11); uncapped optimum has n<=2 chains (informational)\", 1.0, Infinity]]\n"
+     ]
+    }
+   ],
+   "source": [
+    "print()\n",
+    "print(f'{\"identity\":94s} {\"max abs err\":>12s}')\n",
+    "for n_, e_, t_ in CHECKS: print(f'{n_:94s} {e_:12.3e}')\n",
+    "print('CHECKS_JSON:' + json.dumps(CHECKS))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

code/notebooks/0208-statistics-of-evaluation.ipynb

@@ -0,0 +1,687 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "2dc561da",
+   "metadata": {},
+   "source": [
+    "# 0208 -- Statistics of evaluation: exact McNemar, Wilson intervals, cluster bootstrap, power for a +5/244 effect\n",
+    "\n",
+    "Benchmark of `N = 244` problems (the size of miniF2F-test). Identities / claims verified:\n",
+    "1. Exact McNemar two-sided p-value `min(1, 2 P(Bin(n_d,1/2) <= min(b,c)))`: exact rationals = scipy = sign-flip permutation enumeration; **no outcome with `b - c = 5` is significant at 0.05 (best case (5,0): p = 0.0625)**.\n",
+    "2. Wilson interval: closed form = numerical inversion of the score test; exact coverage vs Wald and Clopper-Pearson.\n",
+    "3. Variance decomposition `Var = [Var_x(p) + E p(1-p)/m]/N` and the design effect; cluster-bootstrap vs naive coverage.\n",
+    "4. Paired-difference variance `(Var A + Var B - 2 Cov)/N`.\n",
+    "5. Exact power of McNemar for `delta = 5/244`; normal approximation; MDE; required N.\n",
+    "6. Size of Wald and bootstrap tests with few discordant pairs (both reject (5,0); exact does not)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "c6f8a060",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:32:01.357637Z",
+     "iopub.status.busy": "2026-10-14T09:32:01.357418Z",
+     "iopub.status.idle": "2026-10-14T09:32:02.129875Z",
+     "shell.execute_reply": "2026-10-14T09:32:02.128835Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n",
+    "import math, json, itertools\n",
+    "from fractions import Fraction as Fr\n",
+    "from math import comb\n",
+    "import numpy as np\n",
+    "from scipy import stats\n",
+    "from scipy.optimize import brentq\n",
+    "from scipy.special import betaln\n",
+    "CHECKS = []\n",
+    "def check(name, err, tol):\n",
+    "    CHECKS.append((name, float(err), float(tol)))\n",
+    "    assert err <= tol, (name, float(err), tol)\n",
+    "    print(f'{name:96s} err={float(err):.3e}  (tol {tol:.0e})')\n",
+    "N = 244"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8c953229",
+   "metadata": {},
+   "source": [
+    "## 1. Exact McNemar\n",
+    "Paired binary outcomes per problem: `b` = A solves & B fails, `c` = B solves & A fails. Concordant problems carry no information about the difference\n",
+    "under `H0: P(A=1,B=0) = P(A=0,B=1)`. Conditional on `n_d = b + c`, `b ~ Bin(n_d, 1/2)`.\n",
+    "NAIVE: exact rational binomial tail; independent check 1: scipy; independent check 2: enumerate all `2^{n_d}` sign assignments (permutation test)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "7e8399d5",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:32:02.131822Z",
+     "iopub.status.busy": "2026-10-14T09:32:02.131499Z",
+     "iopub.status.idle": "2026-10-14T09:32:03.388935Z",
+     "shell.execute_reply": "2026-10-14T09:32:03.387984Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1. exact McNemar p: rational tail vs scipy.binomtest (all b,c <= 16)                             err=1.110e-16  (tol 1e-12)\n",
+      "1. exact McNemar p: rational tail vs sign-flip permutation enumeration (n_d <= 14)               err=0.000e+00  (tol 1e-15)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "net gain d=b-c : the discordant counts n_d for which some split is significant (min..max n_d, count)\n",
+      "  d= 3: never significant\n",
+      "  d= 4: never significant\n",
+      "  d= 5: never significant\n",
+      "  d= 6: n_d in [6, 6] (1 values)\n",
+      "  d= 7: n_d in [7, 9] (2 values)\n",
+      "  d= 8: n_d in [8, 12] (3 values)\n",
+      "  d= 9: n_d in [9, 17] (5 values)\n",
+      "  d=10: n_d in [10, 20] (6 values)\n",
+      "  d=11: n_d in [11, 25] (8 values)\n",
+      "  d=12: n_d in [12, 30] (10 values)\n",
+      "p(5,0) = 0.0625   p(6,0) = 0.03125   p(7,1) = 0.0703125   p(8,3) = 0.2265625\n",
+      "1. best possible p for a +5 net gain is (5,0): p = 1/16 > 0.05 (|p - 0.0625|)                    err=0.000e+00  (tol 0e+00)\n",
+      "typical +5/244 outcomes (b,c):  {(5, 0): 0.062, (8, 3): 0.227, (10, 5): 0.302, (15, 10): 0.424, (20, 15): 0.5, (30, 25): 0.59}\n"
+     ]
+    }
+   ],
+   "source": [
+    "def mcnemar_exact_fraction(b, c):\n",
+    "    n = b + c\n",
+    "    if n == 0: return Fr(1)\n",
+    "    k = min(b, c)\n",
+    "    return min(Fr(1), 2 * Fr(sum(comb(n, i) for i in range(k + 1)), 2 ** n))\n",
+    "worst = 0.0; worst_perm = 0.0\n",
+    "for b in range(0, 17):\n",
+    "    for c in range(0, 17):\n",
+    "        n = b + c\n",
+    "        if n == 0: continue\n",
+    "        pf = float(mcnemar_exact_fraction(b, c))\n",
+    "        ps = stats.binomtest(min(b, c), n, 0.5, alternative='two-sided').pvalue\n",
+    "        worst = max(worst, abs(pf - ps))\n",
+    "        if n <= 14:\n",
+    "            signs = np.array(list(itertools.product([-1, 1], repeat=n)))\n",
+    "            cnt = (np.abs(signs.sum(1)) >= abs(b - c)).sum() / len(signs)\n",
+    "            worst_perm = max(worst_perm, abs(pf - cnt))\n",
+    "check('1. exact McNemar p: rational tail vs scipy.binomtest (all b,c <= 16)', worst, 1e-12)\n",
+    "check('1. exact McNemar p: rational tail vs sign-flip permutation enumeration (n_d <= 14)', worst_perm, 1e-15)\n",
+    "# Which net improvements d = b - c can ever be significant at 0.05 (two-sided)?\n",
+    "best = {}\n",
+    "for n in range(1, N + 1):\n",
+    "    for b in range(n + 1):\n",
+    "        c = n - b\n",
+    "        if b - c <= 0: continue\n",
+    "        if float(mcnemar_exact_fraction(b, c)) < 0.05:\n",
+    "            best.setdefault(b - c, []).append(n)\n",
+    "print('net gain d=b-c : the discordant counts n_d for which some split is significant (min..max n_d, count)')\n",
+    "for d in range(3, 13):\n",
+    "    ns_ = best.get(d, [])\n",
+    "    print(f'  d={d:2d}: ' + (f'n_d in [{min(ns_)}, {max(ns_)}] ({len(ns_)} values)' if ns_ else 'never significant'))\n",
+    "for d in (5, 6):\n",
+    "    pass\n",
+    "assert 5 not in best and 4 not in best and 6 in best and min(best[6]) == max(best[6]) == 6\n",
+    "print('p(5,0) =', float(mcnemar_exact_fraction(5, 0)), '  p(6,0) =', float(mcnemar_exact_fraction(6, 0)), '  p(7,1) =', float(mcnemar_exact_fraction(7, 1)), '  p(8,3) =', float(mcnemar_exact_fraction(8, 3)))\n",
+    "check('1. best possible p for a +5 net gain is (5,0): p = 1/16 > 0.05 (|p - 0.0625|)', abs(float(mcnemar_exact_fraction(5, 0)) - 0.0625), 0.0)\n",
+    "print('typical +5/244 outcomes (b,c): ', {(b, b - 5): round(float(mcnemar_exact_fraction(b, b - 5)), 3) for b in (5, 8, 10, 15, 20, 30)})"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d98651ac",
+   "metadata": {},
+   "source": [
+    "## 2. Wilson interval\n",
+    "Invert the score test `|p_hat - p| <= z sqrt(p(1-p)/n)`. Squaring: `(1 + z^2/n) p^2 - (2 p_hat + z^2/n) p + p_hat^2 <= 0`, a quadratic in `p` whose roots are\n",
+    "`p = [p_hat + z^2/(2n) +- z sqrt(p_hat(1-p_hat)/n + z^2/(4n^2))] / (1 + z^2/n)`.\n",
+    "NAIVE: root-find the score-test equation numerically. EFFICIENT: closed form. Then exact coverage over a grid of true `p`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "556d82e0",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:32:03.390688Z",
+     "iopub.status.busy": "2026-10-14T09:32:03.390516Z",
+     "iopub.status.idle": "2026-10-14T09:32:03.531057Z",
+     "shell.execute_reply": "2026-10-14T09:32:03.530078Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2. Wilson closed form vs numerical inversion of the score test (max |endpoint diff|, x=1..243)   err=4.905e-13  (tol 1e-10)\n",
+      "  203/244 = 0.8320  Wilson 95% = [0.7800, 0.8737]  (half-width 0.0468 = 11.4 problems)\n",
+      "  208/244 = 0.8525  Wilson 95% = [0.8025, 0.8915]  (half-width 0.0445 = 10.9 problems)\n",
+      "  213/244 = 0.8730  Wilson 95% = [0.8253, 0.9090]  (half-width 0.0419 = 10.2 problems)\n",
+      "  coverage of nominal-95% wilson: mean 0.9507  min 0.9402  at p=0.83: 0.9495  at p=0.02: 0.9664\n",
+      "  coverage of nominal-95% wald  : mean 0.9436  min 0.8634  at p=0.83: 0.9437  at p=0.02: 0.8634\n",
+      "  coverage of nominal-95% cp    : mean 0.9595  min 0.9507  at p=0.83: 0.9597  at p=0.02: 0.9819\n"
+     ]
+    }
+   ],
+   "source": [
+    "z = stats.norm.ppf(0.975)\n",
+    "def wilson(x, n, z=z):\n",
+    "    ph = x / n; den = 1 + z * z / n\n",
+    "    ctr = (ph + z * z / (2 * n)) / den; half = z / den * math.sqrt(ph * (1 - ph) / n + z * z / (4 * n * n))\n",
+    "    return ctr - half, ctr + half\n",
+    "def wilson_numeric(x, n, z=z):\n",
+    "    ph = x / n\n",
+    "    f = lambda p: (ph - p) ** 2 - z * z * p * (1 - p) / n\n",
+    "    lo = brentq(f, 1e-12, ph) if x > 0 else 0.0\n",
+    "    hi = brentq(f, ph, 1 - 1e-12) if x < n else 1.0\n",
+    "    return lo, hi\n",
+    "worst = max(max(abs(a - b_) for a, b_ in zip(wilson(x, N), wilson_numeric(x, N))) for x in range(1, N))\n",
+    "check('2. Wilson closed form vs numerical inversion of the score test (max |endpoint diff|, x=1..243)', worst, 1e-10)\n",
+    "for x in (203, 208, 213):\n",
+    "    lo, hi = wilson(x, N); print(f'  {x}/{N} = {x/N:.4f}  Wilson 95% = [{lo:.4f}, {hi:.4f}]  (half-width {(hi-lo)/2:.4f} = {(hi-lo)/2*N:.1f} problems)')\n",
+    "def wald(x, n):\n",
+    "    ph = x / n; h = z * math.sqrt(ph * (1 - ph) / n); return ph - h, ph + h\n",
+    "def cp(x, n):\n",
+    "    lo = 0.0 if x == 0 else stats.beta.ppf(0.025, x, n - x + 1); hi = 1.0 if x == n else stats.beta.ppf(0.975, x + 1, n - x)\n",
+    "    return lo, hi\n",
+    "xs = np.arange(N + 1); ci_cache = {nm: np.array([f(x, N) for x in xs]) for nm, f in (('wilson', wilson), ('wald', wald), ('cp', cp))}\n",
+    "def coverage(nm, p):\n",
+    "    pm = stats.binom.pmf(xs, N, p); lo, hi = ci_cache[nm][:, 0], ci_cache[nm][:, 1]\n",
+    "    return float((pm * ((lo <= p) & (p <= hi))).sum())\n",
+    "pgrid = np.arange(0.01, 0.995, 0.005)\n",
+    "cov = {nm: np.array([coverage(nm, p) for p in pgrid]) for nm in ci_cache}\n",
+    "for nm in cov: print(f'  coverage of nominal-95% {nm:6s}: mean {cov[nm].mean():.4f}  min {cov[nm].min():.4f}  at p=0.83: {coverage(nm, 0.83):.4f}  at p=0.02: {coverage(nm, 0.02):.4f}')\n",
+    "assert cov['cp'].min() >= 0.95 - 1e-9 and cov['wald'].min() < 0.93 and abs(cov['wilson'].mean() - 0.95) < 0.01\n",
+    "CHECKS.append(('2. min exact coverage over p in [0.01,0.995]: Wilson / Wald / Clopper-Pearson = %.4f / %.4f / %.4f (informational)' % (cov['wilson'].min(), cov['wald'].min(), cov['cp'].min()), float(cov['wilson'].min()), float('inf')))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c7da2ddd",
+   "metadata": {},
+   "source": [
+    "## 3. Clustering by problem: variance decomposition, design effect, cluster bootstrap\n",
+    "`m` samples per problem; the unit of independence is the **problem**. Super-population: `p_x ~ Beta(1.4, 0.7)`, `c_x | p_x ~ Bin(m, p_x)`, `p_hat_x = c_x / m`, estimator `mean_x p_hat_x`.\n",
+    "Total variance: `Var = [Var(p_x) + E p_x(1-p_x)/m] / N`. A sample-level (naive) SE treats all `N m` outcomes as independent: `p_bar(1-p_bar)/(N m)`.\n",
+    "Design effect `1 + (m-1) rho`, `rho = Var(p_x) / (p_bar(1-p_bar))`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "9430c80f",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:32:03.533001Z",
+     "iopub.status.busy": "2026-10-14T09:32:03.532814Z",
+     "iopub.status.idle": "2026-10-14T09:32:06.152740Z",
+     "shell.execute_reply": "2026-10-14T09:32:06.151923Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "p_bar=0.6667  Var(p_x)=0.0717  ICC rho=0.3226  design effect=5.839  -> naive SE underestimates by x2.416\n",
+      "3. design effect identity Var_total / naive_Var = 1 + (m-1) rho                                  err=8.882e-16  (tol 1e-12)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "empirical Var(estimator)=3.3345e-04   theory=3.3235e-04   E[SE_cluster^2]=3.3236e-04   E[SE_naive^2]=5.6833e-05\n",
+      "3. empirical variance of the estimator vs [Var(p)+E p(1-p)/m]/N (relative)                       err=3.303e-03  (tol 3e-02)\n",
+      "3. cluster SE^2 is unbiased for the variance (relative)                                          err=3.929e-05  (tol 3e-02)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "coverage of nominal 95% CI for pass@1 over 600 replications: cluster bootstrap 0.950, naive sample bootstrap 0.585, Wilson on pooled 3904 outcomes 0.602\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "pass@8 (n=16 per problem, unbiased estimator): true 0.9622, cluster-bootstrap CI coverage 0.927\n"
+     ]
+    }
+   ],
+   "source": [
+    "a, b = 1.4, 0.7; m = 16\n",
+    "p_bar = a / (a + b); var_p = a * b / ((a + b) ** 2 * (a + b + 1)); E_pq = p_bar - (var_p + p_bar ** 2)\n",
+    "var_true = (var_p + E_pq / m) / N\n",
+    "rho = var_p / (p_bar * (1 - p_bar)); deff = 1 + (m - 1) * rho\n",
+    "naive_var = p_bar * (1 - p_bar) / (N * m)\n",
+    "print(f'p_bar={p_bar:.4f}  Var(p_x)={var_p:.4f}  ICC rho={rho:.4f}  design effect={deff:.3f}  -> naive SE underestimates by x{math.sqrt(deff):.3f}')\n",
+    "check('3. design effect identity Var_total / naive_Var = 1 + (m-1) rho', abs(var_true / naive_var - deff), 1e-12)\n",
+    "rs = np.random.RandomState(0); R = 30000\n",
+    "px = rs.beta(a, b, size=(R, N)); cx = rs.binomial(m, px); phx = cx / m\n",
+    "est = phx.mean(1)\n",
+    "se2_cluster = phx.var(1, ddof=1) / N; se2_naive = est * (1 - est) / (N * m)\n",
+    "print(f'empirical Var(estimator)={est.var():.4e}   theory={var_true:.4e}   E[SE_cluster^2]={se2_cluster.mean():.4e}   E[SE_naive^2]={se2_naive.mean():.4e}')\n",
+    "check('3. empirical variance of the estimator vs [Var(p)+E p(1-p)/m]/N (relative)', abs(est.var() / var_true - 1), 0.03)\n",
+    "check('3. cluster SE^2 is unbiased for the variance (relative)', abs(se2_cluster.mean() / var_true - 1), 0.03)\n",
+    "assert se2_naive.mean() < 0.5 * var_true\n",
+    "# coverage of 95% CIs for pass@1: cluster bootstrap vs naive sample-level bootstrap vs Wilson on pooled counts\n",
+    "Rcov, B = 600, 400\n",
+    "cov_cl = cov_nv = cov_wi = 0\n",
+    "rb = np.random.RandomState(1)\n",
+    "for r in range(Rcov):\n",
+    "    px_ = rb.beta(a, b, size=N); cx_ = rb.binomial(m, px_); ph = cx_ / m; e = ph.mean()\n",
+    "    idx = rb.randint(0, N, size=(B, N)); boot = ph[idx].mean(1); lo, hi = np.percentile(boot, [2.5, 97.5]); cov_cl += (lo <= p_bar <= hi)\n",
+    "    nb = rb.binomial(N * m, e, size=B) / (N * m); lo, hi = np.percentile(nb, [2.5, 97.5]); cov_nv += (lo <= p_bar <= hi)\n",
+    "    lo, hi = wilson(int(cx_.sum()), N * m); cov_wi += (lo <= p_bar <= hi)\n",
+    "print(f'coverage of nominal 95% CI for pass@1 over {Rcov} replications: cluster bootstrap {cov_cl/Rcov:.3f}, naive sample bootstrap {cov_nv/Rcov:.3f}, Wilson on pooled {N*m} outcomes {cov_wi/Rcov:.3f}')\n",
+    "assert cov_cl / Rcov > 0.92 and cov_nv / Rcov < 0.75\n",
+    "CHECKS.append(('3. coverage of nominal-95%% CI: cluster bootstrap %.3f, naive %.3f (informational)' % (cov_cl / Rcov, cov_nv / Rcov), float(cov_cl / Rcov), float('inf')))\n",
+    "\n",
+    "# pass@k with the unbiased per-problem estimator, cluster bootstrap\n",
+    "k = 8\n",
+    "def rho_nk(n, c, k): \n",
+    "    out = np.ones_like(c, dtype=float)\n",
+    "    for cc in range(n + 1): out[c == cc] = 1.0 if n - cc < k else 1.0 - float(np.prod(1.0 - k / np.arange(n - cc + 1, n + 1)))\n",
+    "    return out\n",
+    "truth_k = 1 - math.exp(betaln(a, b + k) - betaln(a, b))\n",
+    "cov_k = 0; rb = np.random.RandomState(2)\n",
+    "for r in range(Rcov):\n",
+    "    px_ = rb.beta(a, b, size=N); cx_ = rb.binomial(m, px_); est_x = rho_nk(m, cx_, k)\n",
+    "    idx = rb.randint(0, N, size=(B, N)); boot = est_x[idx].mean(1); lo, hi = np.percentile(boot, [2.5, 97.5]); cov_k += (lo <= truth_k <= hi)\n",
+    "print(f'pass@{k} (n={m} per problem, unbiased estimator): true {truth_k:.4f}, cluster-bootstrap CI coverage {cov_k/Rcov:.3f}')\n",
+    "assert cov_k / Rcov > 0.92"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e6dd44fc",
+   "metadata": {},
+   "source": [
+    "## 4. Paired differences\n",
+    "`Delta_x = hat p^A_x - hat p^B_x`. `Var(mean Delta) = (Var A + Var B - 2 Cov(A,B)) / N` on the problem level. Unpaired analysis ignores the covariance."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "49028f15",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:32:06.154500Z",
+     "iopub.status.busy": "2026-10-14T09:32:06.154336Z",
+     "iopub.status.idle": "2026-10-14T09:32:07.402778Z",
+     "shell.execute_reply": "2026-10-14T09:32:07.401927Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "empirical Var(mean Delta)=9.9803e-05;  paired formula 9.9418e-05;  unpaired (ignores Cov) 7.8556e-04;  ratio unpaired/paired = 7.90\n",
+      "4. paired variance formula vs empirical (relative)                                               err=3.864e-03  (tol 3e-02)\n"
+     ]
+    }
+   ],
+   "source": [
+    "rs = np.random.RandomState(3); R = 20000\n",
+    "lA = rs.randn(R, N) * 1.8 + 0.5; lB = lA + 0.15 + rs.randn(R, N) * 0.4     # B = A shifted, correlated through the problem\n",
+    "pA, pB = 1 / (1 + np.exp(-lA)), 1 / (1 + np.exp(-lB))\n",
+    "cA, cB = rs.binomial(m, pA) / m, rs.binomial(m, pB) / m\n",
+    "D = (cB - cA).mean(1)\n",
+    "var_pair = np.mean(((cB - cA).var(1, ddof=1)) / N)\n",
+    "var_unpair = np.mean((cA.var(1, ddof=1) + cB.var(1, ddof=1)) / N)\n",
+    "print(f'empirical Var(mean Delta)={D.var():.4e};  paired formula {var_pair:.4e};  unpaired (ignores Cov) {var_unpair:.4e};  ratio unpaired/paired = {var_unpair/var_pair:.2f}')\n",
+    "check('4. paired variance formula vs empirical (relative)', abs(var_pair / D.var() - 1), 0.03)\n",
+    "assert var_unpair > 1.5 * var_pair"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ae7bfc19",
+   "metadata": {},
+   "source": [
+    "## 5. Power analysis for a net gain of 5/244\n",
+    "Outcomes per problem fall in 4 cells; the test uses only `(b, c)`. Joint law `(b, c, rest) ~ Multinomial(N; p_b, p_c, 1-p_b-p_c)`, with `delta = p_b - p_c = 5/244` and discordance `d = p_b + p_c`.\n",
+    "EXACT power = `sum_{n_d} Bin(n_d; N, d) sum_b Bin(b; n_d, p_b/d) 1{ exact p(b, n_d - b) < 0.05 }`.\n",
+    "NORMAL approximation (Connor 1987): `N >= [z_{a/2} sqrt(d) + z_b sqrt(d - delta^2)]^2 / delta^2`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "f74c8d64",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:32:07.404574Z",
+     "iopub.status.busy": "2026-10-14T09:32:07.404402Z",
+     "iopub.status.idle": "2026-10-14T09:32:55.741131Z",
+     "shell.execute_reply": "2026-10-14T09:32:55.739908Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "discordance d   exact power at N=244   normal-approx power   N needed (normal, 80%)   MDE at N=244 for 80% power (problems)\n",
+      "    0.03              0.306                 0.455                   558                  7.3\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "    0.06              0.181                 0.256                  1119                 10.9\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "    0.10              0.122                 0.171                  1867                 14.3\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "    0.15              0.097                 0.128                  2801                 17.5\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "    0.20              0.086                 0.106                  3736                 20.1\n",
+      "5. exact power vs normal approximation at d=0.10, N=244 (abs diff; the normal formula ignores discreteness) err=4.952e-02  (tol 1e-01)\n",
+      "5. exact power (d=0.10) vs Monte Carlo over 4e5 benchmarks (abs diff)                            err=2.797e-04  (tol 3e-03)\n",
+      "5. vectorised threshold power = table-based exact power (d=0.10, N=244)                          err=4.663e-15  (tol 1e-12)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "required N for 80% exact power at delta=5/244:  {0.03: (584, 558), 0.06: (1184, 1119), 0.1: (1946, 1867), 0.15: (2887, 2801), 0.2: (3826, 3736)}  (exact, normal approx)\n",
+      "5. required N (exact) vs normal approximation, d=0.10 (relative difference)                      err=4.242e-02  (tol 5e-02)\n"
+     ]
+    }
+   ],
+   "source": [
+    "alpha = 0.05\n",
+    "reject = np.zeros((N + 1, N + 1), dtype=bool)          # reject[n_d, b]\n",
+    "for n_d in range(1, N + 1):\n",
+    "    for bb in range(n_d + 1):\n",
+    "        reject[n_d, bb] = stats.binomtest(min(bb, n_d - bb), n_d, 0.5).pvalue < alpha\n",
+    "def power_exact(N_, d, delta):\n",
+    "    pb = (d + delta) / 2; pc = (d - delta) / 2\n",
+    "    nds = np.arange(0, N_ + 1); w = stats.binom.pmf(nds, N_, d)\n",
+    "    tot = 0.0\n",
+    "    for n_d in range(1, N_ + 1):\n",
+    "        if w[n_d] < 1e-14: continue\n",
+    "        bs = np.arange(0, n_d + 1); tot += w[n_d] * (stats.binom.pmf(bs, n_d, pb / d) * reject[n_d, bs]).sum()\n",
+    "    return float(tot)\n",
+    "za, zb = stats.norm.ppf(1 - alpha / 2), stats.norm.ppf(0.8)\n",
+    "def N_needed_normal(d, delta, power=0.8):\n",
+    "    zb_ = stats.norm.ppf(power); return (za * math.sqrt(d) + zb_ * math.sqrt(d - delta ** 2)) ** 2 / delta ** 2\n",
+    "delta = 5 / N\n",
+    "print('discordance d   exact power at N=244   normal-approx power   N needed (normal, 80%)   MDE at N=244 for 80% power (problems)')\n",
+    "res_pow = {}\n",
+    "for d in (0.03, 0.06, 0.10, 0.15, 0.20):\n",
+    "    pe = power_exact(N, d, delta)\n",
+    "    pn = float(stats.norm.cdf((delta * math.sqrt(N) - za * math.sqrt(d)) / math.sqrt(d - delta ** 2)))\n",
+    "    # minimal delta with exact power >= 0.8\n",
+    "    lo_, hi_ = delta, d - 1e-9\n",
+    "    for _ in range(40):\n",
+    "        mid = (lo_ + hi_) / 2\n",
+    "        if power_exact(N, d, mid) >= 0.8: hi_ = mid\n",
+    "        else: lo_ = mid\n",
+    "    res_pow[d] = (pe, pn, N_needed_normal(d, delta), hi_ * N)\n",
+    "    print(f'   {d:5.2f}            {pe:7.3f}               {pn:7.3f}              {res_pow[d][2]:8.0f}               {hi_*N:6.1f}')\n",
+    "assert all(v[0] < 0.35 for v in res_pow.values())\n",
+    "check('5. exact power vs normal approximation at d=0.10, N=244 (abs diff; the normal formula ignores discreteness)', abs(res_pow[0.10][0] - res_pow[0.10][1]), 0.12)\n",
+    "# Monte Carlo check of the exact-power calculation\n",
+    "rs = np.random.RandomState(5); d = 0.10; pb = (d + delta) / 2; pc = (d - delta) / 2\n",
+    "cnt = rs.multinomial(N, [pb, pc, 1 - pb - pc], size=400_000)\n",
+    "mc_power = reject[cnt[:, 0] + cnt[:, 1], cnt[:, 0]].mean()\n",
+    "check('5. exact power (d=0.10) vs Monte Carlo over 4e5 benchmarks (abs diff)', abs(mc_power - res_pow[0.10][0]), 0.003)\n",
+    "# Vectorised exact power via rejection thresholds: reject iff min(b,c) <= k*(n_d), k* = ppf(alpha/2) - 1 (two-sided exact p = 2 cdf(min) < alpha)\n",
+    "def power_exact_fast(N_, d, delta, alpha=0.05):\n",
+    "    nd = np.arange(1, N_ + 1); w = stats.binom.pmf(nd, N_, d)\n",
+    "    ks = stats.binom.ppf(alpha / 2, nd, 0.5) - 1\n",
+    "    q = (d + delta) / (2 * d)\n",
+    "    ok = ks >= 0\n",
+    "    pr = np.zeros(len(nd)); kk = ks[ok]; nn = nd[ok]\n",
+    "    pr[ok] = stats.binom.cdf(kk, nn, q) + stats.binom.sf(nn - kk - 1, nn, q)\n",
+    "    return float((w * pr).sum())\n",
+    "check('5. vectorised threshold power = table-based exact power (d=0.10, N=244)', abs(power_exact_fast(N, 0.10, delta) - res_pow[0.10][0]), 1e-12)\n",
+    "def N_needed_exact(d, delta, power=0.8):\n",
+    "    lo_, hi_ = 50, 200000\n",
+    "    while hi_ - lo_ > 1:\n",
+    "        mid = (lo_ + hi_) // 2\n",
+    "        # power is not exactly monotone in N because of discreteness: use the first N after which power stays >= target on a window\n",
+    "        if all(power_exact_fast(mid + t, d, delta) >= power for t in range(0, 6)): hi_ = mid\n",
+    "        else: lo_ = mid\n",
+    "    return hi_\n",
+    "n_ex = {d: N_needed_exact(d, delta) for d in (0.03, 0.06, 0.10, 0.15, 0.20)}\n",
+    "print('required N for 80% exact power at delta=5/244: ', {d: (n_ex[d], round(N_needed_normal(d, delta))) for d in n_ex}, ' (exact, normal approx)')\n",
+    "check('5. required N (exact) vs normal approximation, d=0.10 (relative difference)', abs(n_ex[0.10] / N_needed_normal(0.10, delta) - 1), 0.05)\n",
+    "CHECKS.append(('5. N required for 80%% power at delta=5/244, d=0.10: exact %d, normal %.0f (informational)' % (n_ex[0.10], N_needed_normal(0.10, delta)), float(n_ex[0.10]), float('inf')))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cb66649f",
+   "metadata": {},
+   "source": [
+    "## 6. Size of tests with few discordant pairs, and the `(5, 0)` outcome\n",
+    "Under `H0` condition on `n_d` and `b ~ Bin(n_d, 1/2)`. EXACT type-I error of (i) exact McNemar, (ii) Wald `z = (b-c)/sqrt(b+c)` at 1.96, (iii) paired percentile bootstrap."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "91eb7b3f",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:32:55.742903Z",
+     "iopub.status.busy": "2026-10-14T09:32:55.742712Z",
+     "iopub.status.idle": "2026-10-14T09:32:58.063768Z",
+     "shell.execute_reply": "2026-10-14T09:32:58.062659Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "n_d   size(exact)  size(Wald z>1.96)  size(bootstrap 95% CI excludes 0, MC)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "  5      0.0000       0.0625             0.0593\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "  6      0.0313       0.0313             0.0287\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "  8      0.0078       0.0703             0.0627\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      " 10      0.0215       0.0215             0.0307\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      " 12      0.0386       0.0386             0.0373\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      " 15      0.0352       0.0352             0.0373\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      " 20      0.0414       0.0414             0.0440\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      " 30      0.0428       0.0428             0.0420\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      " 60      0.0273       0.0519             0.0533\n",
+      "(b,c)=(5,0): exact p = 0.0625;  Wald z = 2.236 -> p = 0.0253;  bootstrap 2.5% quantile of b* = 1, P(b*=0) = 0.0064\n"
+     ]
+    }
+   ],
+   "source": [
+    "print('n_d   size(exact)  size(Wald z>1.96)  size(bootstrap 95% CI excludes 0, MC)')\n",
+    "rs = np.random.RandomState(6); sizes = {}\n",
+    "for n_d in (5, 6, 8, 10, 12, 15, 20, 30, 60):\n",
+    "    bs = np.arange(0, n_d + 1); pm = stats.binom.pmf(bs, n_d, 0.5)\n",
+    "    s_exact = float((pm * np.array([stats.binomtest(min(x, n_d - x), n_d, 0.5).pvalue < alpha for x in bs])).sum())\n",
+    "    s_wald = float((pm * (np.abs(2 * bs - n_d) / math.sqrt(n_d) >= za)).sum())\n",
+    "    # bootstrap: resample the N problems => (b*, c*, rest*) ~ Multinomial(N, [b, c, rest]/N); CI of (b*-c*) from 400 draws; reject if CI excludes 0\n",
+    "    rej = 0; reps = 1500\n",
+    "    for r in range(reps):\n",
+    "        bb = rs.binomial(n_d, 0.5); cc = n_d - bb\n",
+    "        draws = rs.multinomial(N, [bb / N, cc / N, 1 - n_d / N], size=400); diff = draws[:, 0] - draws[:, 1]\n",
+    "        lo, hi = np.percentile(diff, [2.5, 97.5]); rej += (lo > 0) or (hi < 0)\n",
+    "    s_boot = rej / reps; sizes[n_d] = (s_exact, s_wald, s_boot)\n",
+    "    print(f'{n_d:3d}   {s_exact:9.4f}    {s_wald:9.4f}          {s_boot:9.4f}')\n",
+    "assert all(v[0] <= alpha + 1e-12 for v in sizes.values())\n",
+    "assert sizes[5][1] > alpha and sizes[8][1] > alpha    # Wald exceeds its nominal size at n_d = 5 and 8\n",
+    "CHECKS.append(('6. exact McNemar size <= 0.05 for every n_d tested (max size %.4f)' % max(v[0] for v in sizes.values()), float(max(v[0] for v in sizes.values())), alpha))\n",
+    "z50 = 5 / math.sqrt(5); p_wald = 2 * (1 - stats.norm.cdf(z50))\n",
+    "p0 = (1 - 5 / N) ** N\n",
+    "draws = np.random.RandomState(8).multinomial(N, [5 / N, 0, 1 - 5 / N], size=200000)[:, 0]\n",
+    "print(f'(b,c)=(5,0): exact p = {float(mcnemar_exact_fraction(5,0)):.4f};  Wald z = {z50:.3f} -> p = {p_wald:.4f};  bootstrap 2.5% quantile of b* = {np.percentile(draws, 2.5):.0f}, P(b*=0) = {p0:.4f}')\n",
+    "CHECKS.append(('6. (5,0): exact p = %.4f, Wald p = %.4f (informational: they disagree about significance)' % (0.0625, p_wald), float(p_wald), float('inf')))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "14039858",
+   "metadata": {},
+   "source": [
+    "## Verification table"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "09c2d874",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:32:58.065562Z",
+     "iopub.status.busy": "2026-10-14T09:32:58.065394Z",
+     "iopub.status.idle": "2026-10-14T09:32:58.069524Z",
+     "shell.execute_reply": "2026-10-14T09:32:58.068827Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "identity                                                                                            max abs err\n",
+      "1. exact McNemar p: rational tail vs scipy.binomtest (all b,c <= 16)                                  1.110e-16\n",
+      "1. exact McNemar p: rational tail vs sign-flip permutation enumeration (n_d <= 14)                    0.000e+00\n",
+      "1. best possible p for a +5 net gain is (5,0): p = 1/16 > 0.05 (|p - 0.0625|)                         0.000e+00\n",
+      "2. Wilson closed form vs numerical inversion of the score test (max |endpoint diff|, x=1..243)        4.905e-13\n",
+      "2. min exact coverage over p in [0.01,0.995]: Wilson / Wald / Clopper-Pearson = 0.9402 / 0.8634 / 0.9507 (informational)    9.402e-01\n",
+      "3. design effect identity Var_total / naive_Var = 1 + (m-1) rho                                       8.882e-16\n",
+      "3. empirical variance of the estimator vs [Var(p)+E p(1-p)/m]/N (relative)                            3.303e-03\n",
+      "3. cluster SE^2 is unbiased for the variance (relative)                                               3.929e-05\n",
+      "3. coverage of nominal-95% CI: cluster bootstrap 0.950, naive 0.585 (informational)                   9.500e-01\n",
+      "4. paired variance formula vs empirical (relative)                                                    3.864e-03\n",
+      "5. exact power vs normal approximation at d=0.10, N=244 (abs diff; the normal formula ignores discreteness)    4.952e-02\n",
+      "5. exact power (d=0.10) vs Monte Carlo over 4e5 benchmarks (abs diff)                                 2.797e-04\n",
+      "5. vectorised threshold power = table-based exact power (d=0.10, N=244)                               4.663e-15\n",
+      "5. required N (exact) vs normal approximation, d=0.10 (relative difference)                           4.242e-02\n",
+      "5. N required for 80% power at delta=5/244, d=0.10: exact 1946, normal 1867 (informational)           1.946e+03\n",
+      "6. exact McNemar size <= 0.05 for every n_d tested (max size 0.0428)                                  4.277e-02\n",
+      "6. (5,0): exact p = 0.0625, Wald p = 0.0253 (informational: they disagree about significance)         2.535e-02\n",
+      "CHECKS_JSON:[[\"1. exact McNemar p: rational tail vs scipy.binomtest (all b,c <= 16)\", 1.1102230246251565e-16, 1e-12], [\"1. exact McNemar p: rational tail vs sign-flip permutation enumeration (n_d <= 14)\", 0.0, 1e-15], [\"1. best possible p for a +5 net gain is (5,0): p = 1/16 > 0.05 (|p - 0.0625|)\", 0.0, 0.0], [\"2. Wilson closed form vs numerical inversion of the score test (max |endpoint diff|, x=1..243)\", 4.904965322793942e-13, 1e-10], [\"2. min exact coverage over p in [0.01,0.995]: Wilson / Wald / Clopper-Pearson = 0.9402 / 0.8634 / 0.9507 (informational)\", 0.940211759420751, Infinity], [\"3. design effect identity Var_total / naive_Var = 1 + (m-1) rho\", 8.881784197001252e-16, 1e-12], [\"3. empirical variance of the estimator vs [Var(p)+E p(1-p)/m]/N (relative)\", 0.0033025628714560273, 0.03], [\"3. cluster SE^2 is unbiased for the variance (relative)\", 3.929333853491812e-05, 0.03], [\"3. coverage of nominal-95% CI: cluster bootstrap 0.950, naive 0.585 (informational)\", 0.95, Infinity], [\"4. paired variance formula vs empirical (relative)\", 0.0038637080259513734, 0.03], [\"5. exact power vs normal approximation at d=0.10, N=244 (abs diff; the normal formula ignores discreteness)\", 0.049521981277600394, 0.12], [\"5. exact power (d=0.10) vs Monte Carlo over 4e5 benchmarks (abs diff)\", 0.0002797189205828443, 0.003], [\"5. vectorised threshold power = table-based exact power (d=0.10, N=244)\", 4.6629367034256575e-15, 1e-12], [\"5. required N (exact) vs normal approximation, d=0.10 (relative difference)\", 0.042423303779374866, 0.05], [\"5. N required for 80% power at delta=5/244, d=0.10: exact 1946, normal 1867 (informational)\", 1946.0, Infinity], [\"6. exact McNemar size <= 0.05 for every n_d tested (max size 0.0428)\", 0.04277394525706773, 0.05], [\"6. (5,0): exact p = 0.0625, Wald p = 0.0253 (informational: they disagree about significance)\", 0.0253473186774682, Infinity]]\n"
+     ]
+    }
+   ],
+   "source": [
+    "print()\n",
+    "print(f'{\"identity\":98s} {\"max abs err\":>12s}')\n",
+    "for n_, e_, t_ in CHECKS: print(f'{n_:98s} {e_:12.3e}')\n",
+    "print('CHECKS_JSON:' + json.dumps(CHECKS))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

code/notebooks/0209-variance-components-equivalence.ipynb

@@ -0,0 +1,486 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 0209 - Variance components and equivalence tests for evaluation\n",
+    "Companion to `notes/0209 - Variance Components and Equivalence Tests.pdf`.  CPU only, numpy/scipy, 1 thread.\n",
+    "Running toy: a benchmark of `N = 244` problems, `S` independently trained models (\"seeds\") of one recipe, `m = 16` samples per (seed, problem).\n",
+    "Per-(seed, problem) solve probability `p_sx = expit(mu_x + a_s + b_sx)`, `mu_x ~ N(0.3, 1.5^2)` (problem difficulty), `a_s ~ N(0, 0.25^2)` (seed), `b_sx ~ N(0, 0.5^2)` (seed x problem)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "source": [
+    "import os\n",
+    "os.environ[\"OMP_NUM_THREADS\"] = \"1\"; os.environ[\"OPENBLAS_NUM_THREADS\"] = \"1\"; os.environ[\"MKL_NUM_THREADS\"] = \"1\"\n",
+    "import math, itertools\n",
+    "from fractions import Fraction as Fr\n",
+    "from math import comb\n",
+    "import numpy as np\n",
+    "from scipy import stats\n",
+    "from scipy.special import expit\n",
+    "CHECKS = []\n",
+    "def check(name, err, tol, note=\"\"):\n",
+    "    CHECKS.append((name, float(err), float(tol)))\n",
+    "    assert err <= tol, (name, float(err), tol)\n",
+    "    print(f\"[ok] {name:100s} err={float(err):.3e} (tol {tol:.1e}) {note}\")\n",
+    "N, m = 244, 16\n",
+    "M0, SX, SA, SB = 0.3, 1.5, 0.25, 0.5"
+   ],
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Variance components: problems (A), seeds (B), seed x problem (C), generation (e)\n",
+    "Exchangeable array `p_sx = mu + A_x + B_s + C_sx` (Hoeffding/ANOVA decomposition; components orthogonal).  With outcomes `c_sx ~ Bin(m, p_sx)` and the score of one run `Y_s = (1/(N m)) sum_x c_sx`:\n",
+    "`Var(Y_s) = sA2/N + sB2 + sC2/N + se2/(N m)`,   `se2 = E p(1-p)`.   Only `sB2` survives `N -> inf` and `m -> inf`.\n",
+    "Population components by Gauss-Hermite integration of the logit-normal model (no simulation)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "source": [
+    "gh_x, gh_w = np.polynomial.hermite_e.hermegauss(60); gh_w = gh_w / gh_w.sum()     # standard normal nodes/weights\n",
+    "MU = M0 + SX * gh_x; A_ = SA * gh_x; B_ = SB * gh_x\n",
+    "P3 = expit(MU[:, None, None] + A_[None, :, None] + B_[None, None, :])                # [x, s, b]\n",
+    "W3 = gh_w[:, None, None] * gh_w[None, :, None] * gh_w[None, None, :]\n",
+    "Ep = (P3 * W3).sum(); Ep2 = (P3 ** 2 * W3).sum()\n",
+    "Ax = (P3 * gh_w[None, :, None] * gh_w[None, None, :]).sum((1, 2))                    # E_{s,b} p | x  (weights over s,b sum to 1)\n",
+    "Bs = (P3 * gh_w[:, None, None] * gh_w[None, None, :]).sum((0, 2))                    # E_{x,b} p | s\n",
+    "sA2 = (gh_w * (Ax - Ep) ** 2).sum(); sB2 = (gh_w * (Bs - Ep) ** 2).sum()\n",
+    "sT2 = Ep2 - Ep ** 2; sC2 = sT2 - sA2 - sB2; se2 = Ep - Ep2\n",
+    "print(f\"  E p = {Ep:.4f}; sigma_A^2 = {sA2:.5f}, sigma_B^2 = {sB2:.6f} (sd {math.sqrt(sB2):.4f}), sigma_C^2 = {sC2:.5f}, sigma_e^2 = E p(1-p) = {se2:.5f}\")\n",
+    "def var_Y(S, N=N, m=m): return sA2 / N + sB2 / S + sC2 / (N * S) + se2 / (N * S * m)      # variance of a grand mean over S seeds\n",
+    "print(\"  per-term contribution to Var(Y_s), S=1:  problems %.3e  seed %.3e  seed*problem %.3e  generation %.3e\" % (sA2 / N, sB2, sC2 / N, se2 / (N * m)))\n",
+    "check(\"1. components are non-negative and sum to Var(p)\", abs(sA2 + sB2 + sC2 - sT2) + max(0, -min(sA2, sB2, sC2)), 1e-15)\n",
+    "# simulate replications: each replication redraws problems, seeds, samples\n",
+    "rng = np.random.default_rng(0)\n",
+    "def simulate(R, S, N=N, m=m):\n",
+    "    mu = M0 + SX * rng.standard_normal((R, 1, N)); a = SA * rng.standard_normal((R, S, 1)); b = SB * rng.standard_normal((R, S, N))\n",
+    "    p = expit(mu + a + b)\n",
+    "    c = rng.binomial(m, p)\n",
+    "    return p, c\n",
+    "R1 = 6000\n",
+    "p, c = simulate(R1, 1)\n",
+    "Y = c[:, 0, :].sum(1) / (N * m)\n",
+    "se_rel = math.sqrt(2 / (R1 - 1))\n",
+    "check(\"1. Var(Y_s) over 6000 replications vs sA2/N + sB2 + sC2/N + se2/(Nm)\", abs(Y.var(ddof=1) / var_Y(1) - 1), 4 * se_rel, f\"emp {Y.var(ddof=1):.4e} vs {var_Y(1):.4e}\")\n",
+    "# MoM / ANOVA estimators with S=4 seeds\n",
+    "S = 4; R2 = 4000\n",
+    "p, c = simulate(R2, S)\n",
+    "ybar = c / m                                                                         # cell means [R,S,N]\n",
+    "gm = ybar.mean((1, 2), keepdims=True)\n",
+    "MSx = S * ((ybar.mean(1) - gm[:, :, 0]) ** 2).sum(1) / (N - 1)\n",
+    "MSs = N * ((ybar.mean(2) - gm[:, :, 0]) ** 2).sum(1) / (S - 1)\n",
+    "res = ybar - ybar.mean(1, keepdims=True) - ybar.mean(2, keepdims=True) + gm\n",
+    "MSr = (res ** 2).sum((1, 2)) / ((S - 1) * (N - 1))\n",
+    "pq = (c * (m - c) / (m * (m - 1))).mean((1, 2))                                      # unbiased for E p(1-p)\n",
+    "est = {\"sigma_e^2\": (pq, se2), \"sigma_C^2\": (MSr - pq / m, sC2), \"sigma_A^2\": ((MSx - MSr) / S, sA2), \"sigma_B^2\": ((MSs - MSr) / N, sB2)}\n",
+    "for k, (v, t) in est.items():\n",
+    "    z = (v.mean() - t) / (v.std(ddof=1) / math.sqrt(R2))\n",
+    "    check(f\"1. ANOVA estimator of {k} is unbiased (4000 reps, S=4)\", abs(z), 4.0, f\"mean {v.mean():.6f} vs true {t:.6f}  (z={z:+.2f})\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "  E p = 0.5512; sigma_A^2 = 0.06735, sigma_B^2 = 0.001800 (sd 0.0424), sigma_C^2 = 0.00852, sigma_e^2 = E p(1-p) = 0.16971\n",
+      "  per-term contribution to Var(Y_s), S=1:  problems 2.760e-04  seed 1.800e-03  seed*problem 3.490e-05  generation 4.347e-05\n",
+      "[ok] 1. components are non-negative and sum to Var(p)                                                     err=0.000e+00 (tol 1.0e-15) \n",
+      "[ok] 1. Var(Y_s) over 6000 replications vs sA2/N + sB2 + sC2/N + se2/(Nm)                                 err=2.987e-02 (tol 7.3e-02) emp 2.2191e-03 vs 2.1547e-03\n",
+      "[ok] 1. ANOVA estimator of sigma_e^2 is unbiased (4000 reps, S=4)                                         err=9.497e-01 (tol 4.0e+00) mean 0.169776 vs true 0.169709  (z=+0.95)\n",
+      "[ok] 1. ANOVA estimator of sigma_C^2 is unbiased (4000 reps, S=4)                                         err=1.471e-01 (tol 4.0e+00) mean 0.008513 vs true 0.008515  (z=-0.15)\n",
+      "[ok] 1. ANOVA estimator of sigma_A^2 is unbiased (4000 reps, S=4)                                         err=6.625e-01 (tol 4.0e+00) mean 0.067303 vs true 0.067352  (z=-0.66)\n",
+      "[ok] 1. ANOVA estimator of sigma_B^2 is unbiased (4000 reps, S=4)                                         err=1.567e+00 (tol 4.0e+00) mean 0.001764 vs true 0.001800  (z=-1.57)\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Which term do you pay down with what?  Two recipes, S seeds each, same problems\n",
+    "`delta_hat = Ybar_A - Ybar_B`.  Under H0 (same recipe) the shared problem effect A_x cancels: `Var(delta_hat) = (2/S) [ sB2 + sC2/N + se2/(N m) ]`.\n",
+    "The seed term does not shrink with N or m: a floor on the minimum detectable effect."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "source": [
+    "def var_diff(S, N=N, m=m): return (2 / S) * (sB2 + sC2 / N + se2 / (N * m))\n",
+    "R3 = 6000\n",
+    "for S in (1, 3):\n",
+    "    mu = M0 + SX * rng.standard_normal((R3, 1, N))\n",
+    "    d = []\n",
+    "    for _ in range(2):\n",
+    "        a = SA * rng.standard_normal((R3, S, 1)); b = SB * rng.standard_normal((R3, S, N))\n",
+    "        cc = rng.binomial(m, expit(mu + a + b)); d.append(cc.sum((1, 2)) / (S * N * m))\n",
+    "    dd = d[0] - d[1]\n",
+    "    check(f\"2. Var(delta_hat) for S={S} seeds per recipe vs (2/S)[sB2 + sC2/N + se2/(Nm)]\", abs(dd.var(ddof=1) / var_diff(S) - 1), 4 * math.sqrt(2 / (R3 - 1)), f\"emp {dd.var(ddof=1):.3e} vs {var_diff(S):.3e}\")\n",
+    "z = stats.norm.ppf(0.975) + stats.norm.ppf(0.8)\n",
+    "for S in (1, 3, 5):\n",
+    "    print(f\"  S={S}: MDE(80% power, 5% two-sided) at N=244: {z*math.sqrt(var_diff(S))*100:.2f} pp;  N->inf, m->inf floor: {z*math.sqrt(2*sB2/S)*100:.2f} pp\")\n",
+    "check(\"2. floor is strictly positive: MDE(N=1e9, m=1e9, S=3) equals z sqrt(2 sB2 / 3)\", abs(z * math.sqrt(var_diff(3, N=10**9, m=10**9)) - z * math.sqrt(2 * sB2 / 3)), 1e-6)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "[ok] 2. Var(delta_hat) for S=1 seeds per recipe vs (2/S)[sB2 + sC2/N + se2/(Nm)]                          err=1.212e-02 (tol 7.3e-02) emp 3.803e-03 vs 3.757e-03\n",
+      "[ok] 2. Var(delta_hat) for S=3 seeds per recipe vs (2/S)[sB2 + sC2/N + se2/(Nm)]                          err=1.255e-02 (tol 7.3e-02) emp 1.237e-03 vs 1.252e-03\n",
+      "  S=1: MDE(80% power, 5% two-sided) at N=244: 17.17 pp;  N->inf, m->inf floor: 16.81 pp\n",
+      "  S=3: MDE(80% power, 5% two-sided) at N=244: 9.91 pp;  N->inf, m->inf floor: 9.71 pp\n",
+      "  S=5: MDE(80% power, 5% two-sided) at N=244: 7.68 pp;  N->inf, m->inf floor: 7.52 pp\n",
+      "[ok] 2. floor is strictly positive: MDE(N=1e9, m=1e9, S=3) equals z sqrt(2 sB2 / 3)                       err=2.295e-10 (tol 1.0e-06) \n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Unbiased pass@k from n samples, and its exact variance\n",
+    "`e(n,c,k) = 1 - C(n-c,k)/C(n,k)`; `c ~ Bin(n,p)`  =>  `E e = 1 - (1-p)^k`.  Exact variance by summing over `c = 0..n`.\n",
+    "NAIVE plug-in `1 - (1 - c/n)^k` is biased.  Special cases: `k = 1` gives `p(1-p)/n`;  `k = n` gives `q(1-q)`, `q = 1-(1-p)^n`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "source": [
+    "def e_nck(n, c, k): return 1 - Fr(comb(n - c, k), comb(n, k))\n",
+    "def pk_moments(n, k, p):\n",
+    "    p = Fr(p); pm = [Fr(comb(n, c)) * p ** c * (1 - p) ** (n - c) for c in range(n + 1)]\n",
+    "    E1 = sum(w * e_nck(n, c, k) for c, w in enumerate(pm)); E2 = sum(w * e_nck(n, c, k) ** 2 for c, w in enumerate(pm))\n",
+    "    return E1, E2 - E1 ** 2, sum(w * (1 - (1 - Fr(c, n)) ** k) for c, w in enumerate(pm))\n",
+    "worst = 0; worst_k1 = 0; worst_kn = 0\n",
+    "for n, k, p in [(64, 8, Fr(1, 10)), (64, 32, Fr(1, 50)), (16, 4, Fr(3, 10)), (64, 1, Fr(1, 7)), (64, 64, Fr(1, 100)), (10, 10, Fr(1, 5))]:\n",
+    "    E1, V, Enaive = pk_moments(n, k, p)\n",
+    "    worst = max(worst, abs(float(E1 - (1 - (1 - p) ** k))))\n",
+    "    if k == 1: worst_k1 = max(worst_k1, abs(float(V - p * (1 - p) / n)))\n",
+    "    if k == n: q = 1 - (1 - p) ** n; worst_kn = max(worst_kn, abs(float(V - q * (1 - q))))\n",
+    "check(\"3. E[e(n,c,k)] = 1-(1-p)^k exactly (6 (n,k,p) triples, rational arithmetic)\", worst, 1e-15)\n",
+    "check(\"3. k=1: Var = p(1-p)/n exactly\", worst_k1, 1e-15)\n",
+    "check(\"3. k=n: Var = q(1-q) with q = 1-(1-p)^n exactly\", worst_kn, 1e-15)\n",
+    "print(\"  n=64:        p     k   true pass@k   E[unbiased]   E[naive plug-in] (bias)   sd(unbiased)   sd/true\")\n",
+    "for p in (0.02, 0.1, 0.5):\n",
+    "    for k in (1, 8, 32):\n",
+    "        E1, V, En = pk_moments(64, k, Fr(p).limit_denominator(1000)); g = 1 - (1 - p) ** k\n",
+    "        print(f\"           {p:5.2f} {k:4d}   {g:.5f}      {float(E1):.5f}      {float(En):.5f} ({float(En)-g:+.5f})     {math.sqrt(float(V)):.5f}      {math.sqrt(float(V))/g:.3f}\")\n",
+    "E1, V, En = pk_moments(64, 32, Fr(1, 50))\n",
+    "check(\"3. naive plug-in is biased low at (n=64,k=32,p=.02) by more than 0.07\", 0.0 if (1 - (1 - 0.02) ** 32) - float(En) > 0.07 else 1.0, 0.0, f\"bias {float(En)-(1-0.98**32):+.4f}\")\n",
+    "# MC check of the exact variance (independent implementation via binomial draws) and of the problem-level decomposition\n",
+    "rng = np.random.default_rng(1)\n",
+    "n, k, p0 = 64, 8, 0.1\n",
+    "cs = rng.binomial(n, p0, size=2_000_000)\n",
+    "lg = np.array([float(e_nck(n, c_, k)) for c_ in range(n + 1)])\n",
+    "est = lg[cs]\n",
+    "Vex = float(pk_moments(n, k, Fr(1, 10))[1])\n",
+    "check(\"3. exact Var vs 2e6 Monte Carlo draws (n=64,k=8,p=.1)\", abs(est.var() / Vex - 1), 4 * math.sqrt(2 / 2e6))\n",
+    "# across problems: p_x ~ Beta(1.4, 0.7); Var(mean_x e_x) = (1/N)[Var_x(g(p_x)) + E_x Var(e | p_x)]\n",
+    "a_, b_ = 1.4, 0.7; k = 8; n = 64\n",
+    "xs = rng.beta(a_, b_, size=400_000)                              # for the population quantities\n",
+    "cs = rng.binomial(n, xs)\n",
+    "lgk = np.array([float(e_nck(n, c_, k)) for c_ in range(n + 1)])\n",
+    "es = lgk[cs]; gx = 1 - (1 - xs) ** k\n",
+    "pred = (gx.var() + (es - gx).var()) / N\n",
+    "reps = np.array([lgk[rng.binomial(n, rng.beta(a_, b_, size=N))].mean() for _ in range(6000)])\n",
+    "check(\"3. Var of the benchmark mean pass@8 (N=244, n=64) vs (1/N)[Var_x g + E Var(e|p)]\", abs(reps.var(ddof=1) / pred - 1), 5 * math.sqrt(2 / 5999) + 0.01, f\"emp {reps.var(ddof=1):.3e} vs {pred:.3e}\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "[ok] 3. E[e(n,c,k)] = 1-(1-p)^k exactly (6 (n,k,p) triples, rational arithmetic)                          err=0.000e+00 (tol 1.0e-15) \n",
+      "[ok] 3. k=1: Var = p(1-p)/n exactly                                                                       err=0.000e+00 (tol 1.0e-15) \n",
+      "[ok] 3. k=n: Var = q(1-q) with q = 1-(1-p)^n exactly                                                      err=0.000e+00 (tol 1.0e-15) \n",
+      "  n=64:        p     k   true pass@k   E[unbiased]   E[naive plug-in] (bias)   sd(unbiased)   sd/true\n",
+      "            0.02    1   0.02000      0.02000      0.02000 (-0.00000)     0.01750      0.875\n",
+      "            0.02    8   0.14924      0.14924      0.14185 (-0.00739)     0.12202      0.818\n",
+      "            0.02   32   0.47612      0.47612      0.40018 (-0.07594)     0.32411      0.681\n",
+      "            0.10    1   0.10000      0.10000      0.10000 (+0.00000)     0.03750      0.375\n",
+      "            0.10    8   0.56953      0.56953      0.54892 (-0.02061)     0.14662      0.257\n",
+      "            0.10   32   0.96566      0.96566      0.92934 (-0.03633)     0.07300      0.076\n",
+      "            0.50    1   0.50000      0.50000      0.50000 (+0.00000)     0.06250      0.125\n",
+      "            0.50    8   0.99609      0.99609      0.99418 (-0.00191)     0.00474      0.005\n",
+      "            0.50   32   1.00000      1.00000      1.00000 (-0.00000)     0.00000      0.000\n",
+      "[ok] 3. naive plug-in is biased low at (n=64,k=32,p=.02) by more than 0.07                                err=0.000e+00 (tol 0.0e+00) bias -0.0759\n",
+      "[ok] 3. exact Var vs 2e6 Monte Carlo draws (n=64,k=8,p=.1)                                                err=9.304e-04 (tol 4.0e-03) \n",
+      "[ok] 3. Var of the benchmark mean pass@8 (N=244, n=64) vs (1/N)[Var_x g + E Var(e|p)]                     err=2.299e-02 (tol 1.0e-01) emp 6.460e-05 vs 6.315e-05\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Paired sign-flip test on per-problem differences\n",
+    "`d_x = e^A_x - e^B_x`, statistic `T = sum_x d_x`, null: each `d_x` symmetric about 0 independently  =>  signs exchangeable: p = fraction of the `2^N` sign vectors with `|T'| >= |T|`.\n",
+    "Binary `d_x in {-1,0,1}`: reduces to exact McNemar (note 0208).  Exact enumeration for N<=18; Monte Carlo for N=244; size under the null."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "source": [
+    "def signflip_exact(d):\n",
+    "    d = np.asarray(d, float); n = len(d)\n",
+    "    signs = np.array(list(itertools.product([-1.0, 1.0], repeat=n)))\n",
+    "    T = np.abs(signs @ d); return float((T >= abs(d.sum()) - 1e-12).mean())\n",
+    "def mcnemar(b, c):\n",
+    "    nd = b + c\n",
+    "    return 1.0 if nd == 0 else min(1.0, 2 * sum(comb(nd, i) for i in range(min(b, c) + 1)) / 2 ** nd)\n",
+    "worst = 0\n",
+    "for b_ in range(0, 8):\n",
+    "    for c_ in range(0, 8):\n",
+    "        if b_ + c_ == 0: continue\n",
+    "        d = np.array([1] * b_ + [-1] * c_ + [0] * 3)\n",
+    "        worst = max(worst, abs(signflip_exact(d) - mcnemar(b_, c_)))\n",
+    "check(\"4. sign-flip p on {-1,0,+1} equals exact McNemar for all b,c <= 7 (3 concordant zeros)\", worst, 1e-12)\n",
+    "check(\"4. (5,0): two-sided sign-flip p = 0.0625\", abs(signflip_exact([1] * 5) - 0.0625), 1e-15)\n",
+    "rng = np.random.default_rng(2)\n",
+    "def signflip_mc(d, B=20000):\n",
+    "    d = np.asarray(d, float); s = rng.choice([-1.0, 1.0], size=(B, len(d))); T = np.abs(s @ d)\n",
+    "    return (1 + (T >= abs(d.sum()) - 1e-12).sum()) / (B + 1)\n",
+    "d = rng.standard_normal(14) * 0.2 + 0.05\n",
+    "check(\"4. Monte Carlo p-value (B=20000) vs exact enumeration at N=14\", abs(signflip_mc(d) - signflip_exact(d)), 4 * 0.5 / math.sqrt(20000))\n",
+    "# size under the null: A and B are two draws of the same system; per-problem pass@8 estimators from n=64 samples with problem-specific p_x\n",
+    "reject = 0; R = 1500; N_ = 60\n",
+    "for r in range(R):\n",
+    "    px = rng.beta(a_, b_, size=N_)\n",
+    "    d = lgk[rng.binomial(n, px)] - lgk[rng.binomial(n, px)]\n",
+    "    reject += signflip_mc(d, B=2000) <= 0.05\n",
+    "size = reject / R\n",
+    "check(\"4. size of the sign-flip test under H0 (N=60 problems, shared p_x, 1500 benchmarks)\", abs(size - 0.05) , 4 * math.sqrt(0.05 * 0.95 / R) + 0.01, f\"size {size:.4f}\")\n",
+    "# the null is symmetry, not mean zero: a skewed mean-zero d inflates/deflates size\n",
+    "reject = 0; R = 1500; N_ = 12\n",
+    "for r in range(R):\n",
+    "    d = np.where(rng.random(N_) < 0.1, 9.0, -1.0)           # mean = 0.1*9 - 0.9 = 0, strongly skewed\n",
+    "    reject += signflip_exact(d) <= 0.05 if N_ <= 12 else 0\n",
+    "size_skew = reject / R\n",
+    "print(f\"  skewed mean-zero differences (N=12): size = {size_skew:.4f} (nominal 0.05): the test's null is symmetry\")\n",
+    "check(\"4. skewed mean-zero null: size differs from 0.05 by more than 0.02 (the test is not a test of the mean)\", 0.0 if abs(size_skew - 0.05) > 0.02 else 1.0, 0.0)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "[ok] 4. sign-flip p on {-1,0,+1} equals exact McNemar for all b,c <= 7 (3 concordant zeros)               err=0.000e+00 (tol 1.0e-12) \n",
+      "[ok] 4. (5,0): two-sided sign-flip p = 0.0625                                                             err=0.000e+00 (tol 1.0e-15) \n",
+      "[ok] 4. Monte Carlo p-value (B=20000) vs exact enumeration at N=14                                        err=6.842e-04 (tol 1.4e-02) \n",
+      "[ok] 4. size of the sign-flip test under H0 (N=60 problems, shared p_x, 1500 benchmarks)                  err=8.000e-03 (tol 3.3e-02) size 0.0580\n",
+      "  skewed mean-zero differences (N=12): size = 0.2940 (nominal 0.05): the test's null is symmetry\n",
+      "[ok] 4. skewed mean-zero null: size differs from 0.05 by more than 0.02 (the test is not a test of the mean) err=0.000e+00 (tol 0.0e+00) \n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Fresh seed on the same problems removes generation noise, not selection bias\n",
+    "K = 8 candidate checkpoints with *identical* true population quality theta (null).  On the fixed benchmark each has a quenched problem-specific advantage\n",
+    "`beta_k ~ N(0, sC2/N)` (seed x problem) plus a fresh-generation error `g_k ~ N(0, se2/(N m))`.  Select `k* = argmax (beta_k + g_k)`.\n",
+    "(i) report the same estimate: bias `sigma_tot c_K`;  (ii) re-sample generations on the SAME problems: bias `E beta_{k*} = (sigma_b^2/sigma_tot) c_K`;  (iii) fresh problems: bias 0.\n",
+    "`c_K = E max of K standard normals`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "source": [
+    "sb2 = sC2 / N; sg2 = se2 / (N * m); st2 = sb2 + sg2; K = 8\n",
+    "cK = float(np.mean(np.max(np.random.default_rng(3).standard_normal((2_000_000, K)), 1)))\n",
+    "from scipy import integrate\n",
+    "cK_exact = integrate.quad(lambda x: K * x * stats.norm.pdf(x) * stats.norm.cdf(x) ** (K - 1), -12, 12)[0]\n",
+    "check(\"5. c_8 = E max of 8 N(0,1): quadrature vs 2e6 Monte Carlo\", abs(cK_exact - cK), 5 * 1.0 / math.sqrt(2e6) , f\"c_8={cK_exact:.5f}\")\n",
+    "rng = np.random.default_rng(4); R = 400_000\n",
+    "beta = math.sqrt(sb2) * rng.standard_normal((R, K)); g1 = math.sqrt(sg2) * rng.standard_normal((R, K)); g2 = math.sqrt(sg2) * rng.standard_normal((R, K))\n",
+    "win = np.argmax(beta + g1, axis=1); ix = np.arange(R)\n",
+    "b_same = (beta + g1)[ix, win]; b_regen = (beta + g2)[ix, win]; b_fresh = (math.sqrt(sb2) * rng.standard_normal(R) + math.sqrt(sg2) * rng.standard_normal(R))\n",
+    "th_same, th_regen = math.sqrt(st2) * cK_exact, (sb2 / math.sqrt(st2)) * cK_exact\n",
+    "sd = math.sqrt(st2)\n",
+    "check(\"5(i). same-sample winner estimate bias vs sigma_tot c_K\", abs(b_same.mean() - th_same), 4 * sd / math.sqrt(R), f\"{b_same.mean()*100:.3f} pp vs {th_same*100:.3f} pp\")\n",
+    "check(\"5(ii). fresh generations on the same problems: bias vs (sb2/sigma_tot) c_K\", abs(b_regen.mean() - th_regen), 4 * sd / math.sqrt(R), f\"{b_regen.mean()*100:.3f} pp vs {th_regen*100:.3f} pp  (retained fraction {th_regen/th_same:.3f} = sb2/sigma_tot^2 = {sb2/st2:.3f})\")\n",
+    "check(\"5(iii). fresh problems: bias 0\", abs(b_fresh.mean()), 4 * sd / math.sqrt(R))\n",
+    "check(\"5. retained fraction of the bias after regenerating = sb2/(sb2+sg2) (with m=16: generation is a minor term)\", abs(th_regen / th_same - sb2 / st2), 1e-12)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "[ok] 5. c_8 = E max of 8 N(0,1): quadrature vs 2e6 Monte Carlo                                            err=7.607e-04 (tol 3.5e-03) c_8=1.42360\n",
+      "[ok] 5(i). same-sample winner estimate bias vs sigma_tot c_K                                              err=1.363e-05 (tol 5.6e-05) 1.262 pp vs 1.260 pp\n",
+      "[ok] 5(ii). fresh generations on the same problems: bias vs (sb2/sigma_tot) c_K                           err=5.363e-06 (tol 5.6e-05) 0.561 pp vs 0.561 pp  (retained fraction 0.445 = sb2/sigma_tot^2 = 0.445)\n",
+      "[ok] 5(iii). fresh problems: bias 0                                                                       err=1.369e-05 (tol 5.6e-05) \n",
+      "[ok] 5. retained fraction of the bias after regenerating = sb2/(sb2+sg2) (with m=16: generation is a minor term) err=5.551e-17 (tol 1.0e-12) \n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Equivalence (TOST) versus \"no significant difference\"; minimum detectable effect\n",
+    "`delta_hat ~ N(delta, s^2)`.  Equivalence margin `Delta`.  TOST: reject `H0: |delta| >= Delta` iff both one-sided tests reject: `p_TOST = max(p_L, p_U)`.\n",
+    "Duality: TOST at level alpha rejects iff the (1-2 alpha) CI lies inside `(-Delta, Delta)`.  Power at `delta=0`: `2 Phi(Delta/s - z_{1-alpha}) - 1`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "source": [
+    "al = 0.05; zA = stats.norm.ppf(1 - al)\n",
+    "def tost_p(dh, s, Dl):\n",
+    "    p_upper = stats.norm.cdf((dh - Dl) / s)          # H0: delta >= Delta\n",
+    "    p_lower = stats.norm.sf((dh + Dl) / s)           # H0: delta <= -Delta\n",
+    "    return np.maximum(p_lower, p_upper)\n",
+    "rng = np.random.default_rng(5)\n",
+    "dh = rng.normal(0, 0.03, 100000); s_ = rng.uniform(0.005, 0.03, 100000); Dl = 0.02\n",
+    "rej_p = tost_p(dh, s_, Dl) < al\n",
+    "rej_ci = (dh - zA * s_ > -Dl) & (dh + zA * s_ < Dl)\n",
+    "check(\"6. TOST p<alpha  <=>  90% CI inside (-Delta, Delta) (100000 random cases)\", float((rej_p != rej_ci).sum()), 0.0)\n",
+    "# size at the boundary and power at delta = 0\n",
+    "R = 400_000\n",
+    "for dl in (Dl, -Dl):\n",
+    "    sm = 0.008; x = rng.normal(dl, sm, R); r = (tost_p(x, sm, Dl) < al).mean()\n",
+    "    check(f\"6. size at delta={dl:+.2f} (boundary of H0), s=0.008: TOST rejection prob = alpha\", abs(r - al), 4 * math.sqrt(al * (1 - al) / R), f\"{r:.4f}\")\n",
+    "sm = 0.008; x = rng.normal(0, sm, R); r = (tost_p(x, sm, Dl) < al).mean(); pw = 2 * stats.norm.cdf(Dl / sm - zA) - 1\n",
+    "check(\"6. power at delta=0: 2 Phi(Delta/s - z) - 1 vs MC\", abs(r - pw), 4 * math.sqrt(0.25 / R), f\"{r:.4f} vs {pw:.4f}\")\n",
+    "# difference test vs equivalence at a true effect larger than the margin: N=244, discordance d=0.1\n",
+    "dpair = 0.1; sd_pair = math.sqrt(dpair); s244 = sd_pair / math.sqrt(244)\n",
+    "true_delta = 0.03\n",
+    "x = rng.normal(true_delta, s244, R)\n",
+    "p_diff = 2 * stats.norm.sf(np.abs(x) / s244)\n",
+    "miss = (p_diff >= 0.05).mean(); fake_equiv = (tost_p(x, s244, Dl) < al).mean()\n",
+    "print(f\"  N=244, d=0.10 (s = {s244:.4f}), true delta=+0.03 > Delta=0.02:  difference test misses it {miss*100:.1f}% of the time; TOST declares equivalence {fake_equiv*100:.2f}%\")\n",
+    "check(\"6. 'not significant' is not equivalence: with a true effect beyond the margin, difference test fails to reject often, TOST seldom 'accepts' (<= alpha)\", 0.0 if (miss > 0.4 and fake_equiv <= al + 4 * math.sqrt(al * 0.95 / R)) else 1.0, 0.0)\n",
+    "# MDE (paired, one run per benchmark) and required N, cross-checked against note 0208's table (normal approximation, d = 0.10, delta = 5/244)\n",
+    "Nn = 244; mde = (stats.norm.ppf(0.975) + stats.norm.ppf(0.8)) * sd_pair / math.sqrt(Nn)\n",
+    "x = rng.normal(mde, s244, R); pw = (2 * stats.norm.sf(np.abs(x) / s244) < 0.05).mean()\n",
+    "check(\"6. MDE = (z_{.975}+z_{.8}) sigma_d/sqrt(N) gives 80% power (Gaussian paired mean)\", abs(pw - 0.8), 4 * math.sqrt(0.16 / R) + 0.0035, f\"MDE = {mde*Nn:.2f} problems, power {pw:.4f}  (0208 exact trinomial: 14.3)\")\n",
+    "dlt = 5 / 244; d_ = 0.10\n",
+    "N_norm = ((stats.norm.ppf(0.975) * math.sqrt(d_) + stats.norm.ppf(0.8) * math.sqrt(d_ - dlt ** 2)) / dlt) ** 2\n",
+    "check(\"6. N for 80% power at delta=5/244, d=.10 equals note 0208's normal-approximation value 1867\", abs(round(N_norm) - 1867), 1, f\"N={N_norm:.1f}\")\n",
+    "Req = (( stats.norm.ppf(1 - al) + stats.norm.ppf(0.9)) * sd_pair / Dl) ** 2        # TOST 80% power at delta=0:  Delta/s - z = z_{0.9}\n",
+    "pw_req = 2 * stats.norm.cdf(Dl / (sd_pair / math.sqrt(Req)) - zA) - 1\n",
+    "check(\"6. N for 80% TOST power at delta=0, Delta=0.02, d=.10: closed form reproduces 80% power\", abs(pw_req - 0.8), 1e-9, f\"N = {Req:.0f}\")\n",
+    "print(f\"  (equivalence within 2 points needs N ~ {Req:.0f} problems at 10% discordance; detecting a 5-problem gain needs ~1867)\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "[ok] 6. TOST p<alpha  <=>  90% CI inside (-Delta, Delta) (100000 random cases)                            err=0.000e+00 (tol 0.0e+00) \n",
+      "[ok] 6. size at delta=+0.02 (boundary of H0), s=0.008: TOST rejection prob = alpha                        err=4.675e-04 (tol 1.4e-03) 0.0495\n",
+      "[ok] 6. size at delta=-0.02 (boundary of H0), s=0.008: TOST rejection prob = alpha                        err=1.425e-04 (tol 1.4e-03) 0.0499\n",
+      "[ok] 6. power at delta=0: 2 Phi(Delta/s - z) - 1 vs MC                                                    err=7.526e-04 (tol 3.2e-03) 0.6083 vs 0.6075\n",
+      "  N=244, d=0.10 (s = 0.0202), true delta=+0.03 > Delta=0.02:  difference test misses it 68.4% of the time; TOST declares equivalence 0.00%\n",
+      "[ok] 6. 'not significant' is not equivalence: with a true effect beyond the margin, difference test fails to reject often, TOST seldom 'accepts' (<= alpha) err=0.000e+00 (tol 0.0e+00) \n",
+      "[ok] 6. MDE = (z_{.975}+z_{.8}) sigma_d/sqrt(N) gives 80% power (Gaussian paired mean)                    err=5.500e-04 (tol 6.0e-03) MDE = 13.84 problems, power 0.7994  (0208 exact trinomial: 14.3)\n",
+      "[ok] 6. N for 80% power at delta=5/244, d=.10 equals note 0208's normal-approximation value 1867          err=0.000e+00 (tol 1.0e+00) N=1866.8\n",
+      "[ok] 6. N for 80% TOST power at delta=0, Delta=0.02, d=.10: closed form reproduces 80% power              err=2.220e-16 (tol 1.0e-09) N = 2141\n",
+      "  (equivalence within 2 points needs N ~ 2141 problems at 10% discordance; detecting a 5-problem gain needs ~1867)\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Summary"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "source": [
+    "for n_, e_, t_ in CHECKS: print(f\"PASS {n_[:105]}\")\n",
+    "print(f\"\\n{len(CHECKS)} checks, all passed\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "PASS 1. components are non-negative and sum to Var(p)\n",
+      "PASS 1. Var(Y_s) over 6000 replications vs sA2/N + sB2 + sC2/N + se2/(Nm)\n",
+      "PASS 1. ANOVA estimator of sigma_e^2 is unbiased (4000 reps, S=4)\n",
+      "PASS 1. ANOVA estimator of sigma_C^2 is unbiased (4000 reps, S=4)\n",
+      "PASS 1. ANOVA estimator of sigma_A^2 is unbiased (4000 reps, S=4)\n",
+      "PASS 1. ANOVA estimator of sigma_B^2 is unbiased (4000 reps, S=4)\n",
+      "PASS 2. Var(delta_hat) for S=1 seeds per recipe vs (2/S)[sB2 + sC2/N + se2/(Nm)]\n",
+      "PASS 2. Var(delta_hat) for S=3 seeds per recipe vs (2/S)[sB2 + sC2/N + se2/(Nm)]\n",
+      "PASS 2. floor is strictly positive: MDE(N=1e9, m=1e9, S=3) equals z sqrt(2 sB2 / 3)\n",
+      "PASS 3. E[e(n,c,k)] = 1-(1-p)^k exactly (6 (n,k,p) triples, rational arithmetic)\n",
+      "PASS 3. k=1: Var = p(1-p)/n exactly\n",
+      "PASS 3. k=n: Var = q(1-q) with q = 1-(1-p)^n exactly\n",
+      "PASS 3. naive plug-in is biased low at (n=64,k=32,p=.02) by more than 0.07\n",
+      "PASS 3. exact Var vs 2e6 Monte Carlo draws (n=64,k=8,p=.1)\n",
+      "PASS 3. Var of the benchmark mean pass@8 (N=244, n=64) vs (1/N)[Var_x g + E Var(e|p)]\n",
+      "PASS 4. sign-flip p on {-1,0,+1} equals exact McNemar for all b,c <= 7 (3 concordant zeros)\n",
+      "PASS 4. (5,0): two-sided sign-flip p = 0.0625\n",
+      "PASS 4. Monte Carlo p-value (B=20000) vs exact enumeration at N=14\n",
+      "PASS 4. size of the sign-flip test under H0 (N=60 problems, shared p_x, 1500 benchmarks)\n",
+      "PASS 4. skewed mean-zero null: size differs from 0.05 by more than 0.02 (the test is not a test of the mean)\n",
+      "PASS 5. c_8 = E max of 8 N(0,1): quadrature vs 2e6 Monte Carlo\n",
+      "PASS 5(i). same-sample winner estimate bias vs sigma_tot c_K\n",
+      "PASS 5(ii). fresh generations on the same problems: bias vs (sb2/sigma_tot) c_K\n",
+      "PASS 5(iii). fresh problems: bias 0\n",
+      "PASS 5. retained fraction of the bias after regenerating = sb2/(sb2+sg2) (with m=16: generation is a minor ter\n",
+      "PASS 6. TOST p<alpha  <=>  90% CI inside (-Delta, Delta) (100000 random cases)\n",
+      "PASS 6. size at delta=+0.02 (boundary of H0), s=0.008: TOST rejection prob = alpha\n",
+      "PASS 6. size at delta=-0.02 (boundary of H0), s=0.008: TOST rejection prob = alpha\n",
+      "PASS 6. power at delta=0: 2 Phi(Delta/s - z) - 1 vs MC\n",
+      "PASS 6. 'not significant' is not equivalence: with a true effect beyond the margin, difference test fails to r\n",
+      "PASS 6. MDE = (z_{.975}+z_{.8}) sigma_d/sqrt(N) gives 80% power (Gaussian paired mean)\n",
+      "PASS 6. N for 80% power at delta=5/244, d=.10 equals note 0208's normal-approximation value 1867\n",
+      "PASS 6. N for 80% TOST power at delta=0, Delta=0.02, d=.10: closed form reproduces 80% power\n",
+      "\n",
+      "33 checks, all passed\n"
+     ]
+    }
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

code/notebooks/0210-data-selection-contamination.ipynb

@@ -0,0 +1,500 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 0210 - Data selection and contamination\n",
+    "Companion to `notes/0210 - Data Selection and Contamination.pdf`.  CPU only, numpy/scipy, 1 thread.\n",
+    "1. classifier-based harvesting (fastText-style) as a selection problem;  2. pass-rate band selection as importance weighting;\n",
+    "3. regression to the mean when the same samples select and evaluate;  4. MinHash / Jaccard estimation error, LSH, containment;  5. membership inference from NLL gaps."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "source": [
+    "import os\n",
+    "os.environ[\"OMP_NUM_THREADS\"] = \"1\"; os.environ[\"OPENBLAS_NUM_THREADS\"] = \"1\"\n",
+    "import math, itertools\n",
+    "from math import comb\n",
+    "import numpy as np\n",
+    "from scipy import stats, special\n",
+    "CHECKS = []\n",
+    "def check(name, err, tol, note=\"\"):\n",
+    "    CHECKS.append((name, float(err), float(tol)))\n",
+    "    assert err <= tol, (name, float(err), tol)\n",
+    "    print(f\"[ok] {name:104s} err={float(err):.3e} (tol {tol:.1e}) {note}\")\n",
+    "Phibar = stats.norm.sf"
+   ],
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. A harvesting classifier is a detector; what you get depends on prevalence and on what a wrong page costs\n",
+    "Pool of `M` pages, a fraction `pi` relevant (\"math\").  Score `s ~ N(d', 1)` for relevant, `N(0, 1)` otherwise (binormal; AUC = Phi(d'/sqrt 2)).\n",
+    "Threshold `t`: TPR = Phibar(t - d'), FPR = Phibar(t).  Precision = pi TPR / (pi TPR + (1-pi) FPR) -- depends on `pi`, not only on AUC.\n",
+    "Downstream value (assumed, to be *measured* by a training run): each selected relevant page gives `+g`, each irrelevant `-h` (dilution).\n",
+    "Net gain `G(t) = M [pi TPR g - (1-pi) FPR h]`.  `G'(t) = 0` where the posterior `P(rel | s=t) = h/(g+h)`: a posterior cutoff, so `t* = d'/2 + ln[(1-pi) h /(pi g)]/d'`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "source": [
+    "rng = np.random.default_rng(0)\n",
+    "dprime, pi, M = 2.5, 0.01, 4_000_000\n",
+    "rel = rng.random(M) < pi\n",
+    "s = rng.standard_normal(M) + dprime * rel\n",
+    "t = 2.0\n",
+    "sel = s >= t\n",
+    "prec_mc = rel[sel].mean(); rec_mc = sel[rel].mean()\n",
+    "tpr, fpr = Phibar(t - dprime), Phibar(t)\n",
+    "prec_th = pi * tpr / (pi * tpr + (1 - pi) * fpr)\n",
+    "check(\"1. precision closed form vs Monte Carlo (4e6 pages)\", abs(prec_mc - prec_th), 4 * math.sqrt(prec_th * (1 - prec_th) / sel.sum()), f\"prec {prec_th:.4f}, recall {tpr:.4f}\")\n",
+    "check(\"1. recall closed form vs Monte Carlo\", abs(rec_mc - tpr), 4 * math.sqrt(tpr * (1 - tpr) / rel.sum()))\n",
+    "auc = stats.norm.cdf(dprime / math.sqrt(2))\n",
+    "print(f\"  AUC={auc:.4f} (d'={dprime}); at a threshold with recall {tpr:.2f}: precision {prec_th:.3f} for pi={pi}; for pi=0.2 the same classifier gives {0.2*tpr/(0.2*tpr+0.8*fpr):.3f}\")\n",
+    "check(\"1. same classifier, same threshold: precision at pi=0.2 exceeds precision at pi=0.01 by > 0.4\", 0.0 if (0.2 * tpr / (0.2 * tpr + 0.8 * fpr)) - prec_th > 0.4 else 1.0, 0.0)\n",
+    "g_, h_ = 1.0, 1.0\n",
+    "ts = np.linspace(0, 6, 60001)\n",
+    "G = pi * Phibar(ts - dprime) * g_ - (1 - pi) * Phibar(ts) * h_\n",
+    "t_star = dprime / 2 + math.log((1 - pi) * h_ / (pi * g_)) / dprime\n",
+    "check(\"1. argmax of the net-gain curve vs closed-form posterior-cutoff threshold\", abs(ts[G.argmax()] - t_star), 2e-4, f\"t*={t_star:.4f}\")\n",
+    "post = lambda t: pi * stats.norm.pdf(t - dprime) / (pi * stats.norm.pdf(t - dprime) + (1 - pi) * stats.norm.pdf(t))\n",
+    "check(\"1. posterior P(rel | s=t*) = h/(g+h)\", abs(post(t_star) - h_ / (g_ + h_)), 1e-12)\n",
+    "def f1(t):\n",
+    "    tp = pi * Phibar(t - dprime); fp = (1 - pi) * Phibar(t); fn = pi - tp\n",
+    "    return 2 * tp / (2 * tp + fp + fn)\n",
+    "t_f1 = ts[np.array([f1(x) for x in ts]).argmax()]\n",
+    "Gof = lambda t: pi * Phibar(t - dprime) * g_ - (1 - pi) * Phibar(t) * h_\n",
+    "print(f\"  max-F1 threshold {t_f1:.3f} vs max-gain threshold {t_star:.3f};  gain/M at F1-opt {Gof(t_f1):.3e} vs gain-opt {Gof(t_star):.3e}\")\n",
+    "check(\"1. the F1-optimal threshold is not the gain-optimal one (differ by > 0.2) and loses downstream gain\", 0.0 if (abs(t_f1 - t_star) > 0.2 and Gof(t_f1) < Gof(t_star)) else 1.0, 0.0)\n",
+    "# seed coverage and iteration: positives split into 'covered' (d_in) and 'uncovered' (d_out) clusters; each round adds a fraction gamma of uncovered clusters to the seed\n",
+    "d_in, d_out, gamma, tt = 3.0, 0.5, 0.5, 2.0\n",
+    "rec = []; cov = 0.3\n",
+    "for r in range(1, 6):\n",
+    "    rec.append(cov * Phibar(tt - d_in) + (1 - cov) * Phibar(tt - d_out)); cov = cov + gamma * (1 - cov)\n",
+    "inc = np.diff(rec)\n",
+    "print(\"  recall by round:\", np.round(rec, 4), \" increments:\", np.round(inc, 4), \" ratio of successive increments:\", np.round(inc[1:] / inc[:-1], 3))\n",
+    "check(\"1. iterative re-seeding: recall rises monotonically and increments shrink geometrically with ratio (1-gamma)\", max(abs(inc[1:] / inc[:-1] - (1 - gamma))), 1e-12)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "[ok] 1. precision closed form vs Monte Carlo (4e6 pages)                                                      err=1.455e-04 (tol 4.9e-03) prec 0.2349, recall 0.6915\n",
+      "[ok] 1. recall closed form vs Monte Carlo                                                                     err=2.719e-03 (tol 9.3e-03) \n",
+      "  AUC=0.9615 (d'=2.5); at a threshold with recall 0.69: precision 0.235 for pi=0.01; for pi=0.2 the same classifier gives 0.884\n",
+      "[ok] 1. same classifier, same threshold: precision at pi=0.2 exceeds precision at pi=0.01 by > 0.4            err=0.000e+00 (tol 0.0e+00) \n",
+      "[ok] 1. argmax of the net-gain curve vs closed-form posterior-cutoff threshold                                err=4.794e-05 (tol 2.0e-04) t*=3.0880\n",
+      "[ok] 1. posterior P(rel | s=t*) = h/(g+h)                                                                     err=1.110e-16 (tol 1.0e-12) \n",
+      "  max-F1 threshold 2.626 vs max-gain threshold 3.088;  gain/M at F1-opt 2.230e-04 vs gain-opt 1.785e-03\n",
+      "[ok] 1. the F1-optimal threshold is not the gain-optimal one (differ by > 0.2) and loses downstream gain      err=0.000e+00 (tol 0.0e+00) \n",
+      "  recall by round: [0.2992 0.5703 0.7058 0.7736 0.8075]  increments: [0.2711 0.1355 0.0678 0.0339]  ratio of successive increments: [0.5 0.5 0.5]\n",
+      "[ok] 1. iterative re-seeding: recall rises monotonically and increments shrink geometrically with ratio (1-gamma) err=6.661e-16 (tol 1.0e-12) \n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Pass-rate band selection is importance weighting of the policy gradient\n",
+    "Bernoulli problem with logit parameter, `n` samples, `c` successes, `p_hat = c/n`.  GRPO-style mean baseline gives per-problem gradient magnitude `ghat = p_hat (1 - p_hat)` (note 0206), `E ghat = (1 - 1/n) p(1-p)`.\n",
+    "Keeping the problem iff `p_hat` lies in a band `[lo, hi]` that is computed from the SAME samples multiplies the gradient by `1{p_hat in band}`:\n",
+    "effective weight on a problem of true pass rate p:  `w_same(p) = E[ghat 1{band}] / (p(1-p))`.  With an independent pilot estimate `p_hat'` instead: `w_ind(p) = (1 - 1/n) P[p_hat' in band]`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "source": [
+    "def bin_pmf(n, p): return np.array([comb(n, c) * p ** c * (1 - p) ** (n - c) for c in range(n + 1)])\n",
+    "n = 8; lo, hi = 1 / 8, 6 / 8\n",
+    "cs = np.arange(n + 1); ph = cs / n; inb = ((ph >= lo) & (ph <= hi)).astype(float); ghat = ph * (1 - ph)\n",
+    "rows = []\n",
+    "for p in (0.02, 0.1, 0.3, 0.5, 0.7, 0.9, 0.98):\n",
+    "    pm = bin_pmf(n, p)\n",
+    "    w_same = (pm * ghat * inb).sum() / (p * (1 - p)); w_ind = (1 - 1 / n) * (pm * inb).sum()\n",
+    "    check_unb = abs((pm * ghat).sum() - (1 - 1 / n) * p * (1 - p))\n",
+    "    assert check_unb < 1e-15\n",
+    "    rows.append((p, w_same, w_ind)); print(f\"  p={p:4.2f}:  w_same={w_same:.4f}   w_ind={w_ind:.4f}   ratio={w_same/w_ind:.3f}\")\n",
+    "check(\"2. E[ghat] = (1-1/n) p(1-p) exactly (7 values of p)\", 0.0, 0.0)\n",
+    "check(\"2. same-sample selection weights exceed independent-pilot weights for every p in the table (selection on the statistic that defines the gradient)\", 0.0 if all(w1 > w2 for _, w1, w2 in rows) else 1.0, 0.0)\n",
+    "# MC check of w_same by simulating samples\n",
+    "p = 0.1; cs_mc = rng.binomial(n, p, size=2_000_000); phm = cs_mc / n\n",
+    "est = (phm * (1 - phm) * ((phm >= lo) & (phm <= hi))).mean() / (p * (1 - p)); th = [r for r in rows if r[0] == p][0][1]\n",
+    "check(\"2. w_same(p=0.1) exact vs 2e6 Monte Carlo\", abs(est / th - 1), 0.01)\n",
+    "# which problems does the band actually admit?  contamination of the band by noise in p_hat (pool p ~ Beta(0.6,1.4), n = 8)\n",
+    "a_, b_ = 0.6, 1.4; band_p = (0.125, 0.75)\n",
+    "px = rng.beta(a_, b_, size=1_000_000); cx = rng.binomial(n, px); sel_ = (cx / n >= lo) & (cx / n <= hi)\n",
+    "frac_out = ((px[sel_] < band_p[0] - 1e-12) | (px[sel_] > band_p[1] + 1e-12)).mean()\n",
+    "# exact via beta-binomial\n",
+    "pmf_bb = np.array([comb(n, c) * special.beta(c + a_, n - c + b_) / special.beta(a_, b_) for c in range(n + 1)])\n",
+    "print(f\"  fraction selected: {sel_.mean():.4f} (exact {pmf_bb[(cs>=1)&(cs<=6)].sum():.4f});  of the selected, true p outside [{band_p[0]},{band_p[1]}]: {frac_out:.3f}\")\n",
+    "check(\"2. beta-binomial exact P[band] vs MC\", abs(sel_.mean() - pmf_bb[(cs >= 1) & (cs <= 6)].sum()), 4 * math.sqrt(0.25 / 1e6))\n",
+    "check(\"2. a substantial share of selected problems lie outside the band in TRUE pass rate (> 15%)\", 0.0 if frac_out > 0.15 else 1.0, 0.0, f\"{frac_out:.3f}\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "  p=0.02:  w_same=0.8750   w_ind=0.1306   ratio=6.701\n",
+      "  p=0.10:  w_same=0.8750   w_ind=0.4983   ratio=1.756\n",
+      "  p=0.30:  w_same=0.8744   w_ind=0.8234   ratio=1.062\n",
+      "  p=0.50:  w_same=0.8613   w_ind=0.8408   ratio=1.024\n",
+      "  p=0.70:  w_same=0.7721   w_ind=0.6516   ratio=1.185\n",
+      "  p=0.90:  w_same=0.4100   w_ind=0.1635   ratio=2.507\n",
+      "  p=0.98:  w_same=0.0999   w_ind=0.0090   ratio=11.044\n",
+      "[ok] 2. E[ghat] = (1-1/n) p(1-p) exactly (7 values of p)                                                      err=0.000e+00 (tol 0.0e+00) \n",
+      "[ok] 2. same-sample selection weights exceed independent-pilot weights for every p in the table (selection on the statistic that defines the gradient) err=0.000e+00 (tol 0.0e+00) \n",
+      "[ok] 2. w_same(p=0.1) exact vs 2e6 Monte Carlo                                                                err=6.072e-04 (tol 1.0e-02) \n",
+      "  fraction selected: 0.6231 (exact 0.6233);  of the selected, true p outside [0.125,0.75]: 0.209\n",
+      "[ok] 2. beta-binomial exact P[band] vs MC                                                                     err=2.182e-04 (tol 2.0e-03) \n",
+      "[ok] 2. a substantial share of selected problems lie outside the band in TRUE pass rate (> 15%)               err=0.000e+00 (tol 0.0e+00) 0.209\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Regression to the mean: the same samples cannot select and evaluate\n",
+    "Pool `p ~ Beta(a, b)`, `c | p ~ Bin(n, p)`.  Posterior `p | c ~ Beta(a + c, b + n - c)`, mean `(a + c)/(a + b + n)`.\n",
+    "Select the \"unsolved\" problems `c = 0`: observed rate 0, true mean `a/(a+b+n)` -- re-sampling them WITHOUT any training shows an apparent gain."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "source": [
+    "a_, b_, n = 0.6, 1.4, 8\n",
+    "post_mean = lambda c: (a_ + c) / (a_ + b_ + n)\n",
+    "px = rng.beta(a_, b_, size=3_000_000); c1 = rng.binomial(n, px); c2 = rng.binomial(n, px)       # c2: fresh samples, NO training between\n",
+    "for cc in (0, 1, 2):\n",
+    "    selc = c1 == cc\n",
+    "    emp = c2[selc].mean() / n; th = post_mean(cc)\n",
+    "    check(f\"3. E[fresh pass rate | first-pass count c={cc}] = (a+c)/(a+b+n)\", abs(emp - th), 4 * math.sqrt(th * (1 - th) / (selc.sum() * n)), f\"observed first-pass rate {cc/n:.3f} -> fresh {emp:.4f} (theory {th:.4f}); apparent 'gain' {emp - cc/n:+.4f} with no training\")\n",
+    "# difference-in-differences with an untrained control from the same selection\n",
+    "tau = 0.03; sel0 = c1 == 0\n",
+    "idx = np.flatnonzero(sel0); half = rng.random(idx.size) < 0.5           # random split of the selected problems into trained / untrained\n",
+    "p_after = np.where(half, np.minimum(px[idx] + tau * (1 - px[idx]), 1.0), px[idx])      # training raises p by tau (1-p) on the trained half only\n",
+    "c_after = rng.binomial(n, p_after) / n\n",
+    "naive = c_after[half].mean() - 0.0; control = c_after[~half].mean() - 0.0\n",
+    "did = naive - control\n",
+    "tau_eff = (tau * (1 - px[idx[half]])).mean()\n",
+    "se = math.sqrt(c_after[half].var() / half.sum() / n + c_after[~half].var() / (~half).sum() / n) + 1e-9\n",
+    "print(f\"  trained-half before/after = {naive:+.4f} (true effect {tau_eff:+.4f}); untrained control 'gain' = {control:+.4f}; DiD = {did:+.4f}\")\n",
+    "check(\"3. naive before/after overstates the effect (regression to the mean) and DiD recovers it\", abs(did - tau_eff), 0.003, f\"naive {naive:+.4f}, DiD {did:+.4f}, true {tau_eff:+.4f}\")\n",
+    "check(\"3. naive before/after is inflated by RTM: excess over truth exceeds 0.03\", 0.0 if (naive - tau_eff) > 0.03 else 1.0, 0.0)\n",
+    "# split-sample selection: select with set 1, evaluate with independent set 2 -> E[fresh | selected] is what an untrained re-evaluation gives (no gain)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "[ok] 3. E[fresh pass rate | first-pass count c=0] = (a+c)/(a+b+n)                                             err=8.957e-05 (tol 3.6e-04) observed first-pass rate 0.000 -> fresh 0.0599 (theory 0.0600); apparent 'gain' +0.0599 with no training\n",
+      "[ok] 3. E[fresh pass rate | first-pass count c=1] = (a+c)/(a+b+n)                                             err=3.516e-04 (tol 7.3e-04) observed first-pass rate 0.125 -> fresh 0.1604 (theory 0.1600); apparent 'gain' +0.0354 with no training\n",
+      "[ok] 3. E[fresh pass rate | first-pass count c=2] = (a+c)/(a+b+n)                                             err=1.404e-04 (tol 1.0e-03) observed first-pass rate 0.250 -> fresh 0.2599 (theory 0.2600); apparent 'gain' +0.0099 with no training\n",
+      "  trained-half before/after = +0.0882 (true effect +0.0282); untrained control 'gain' = +0.0600; DiD = +0.0282\n",
+      "[ok] 3. naive before/after overstates the effect (regression to the mean) and DiD recovers it                 err=9.032e-06 (tol 3.0e-03) naive +0.0882, DiD +0.0282, true +0.0282\n",
+      "[ok] 3. naive before/after is inflated by RTM: excess over truth exceeds 0.03                                 err=0.000e+00 (tol 0.0e+00) \n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. MinHash: unbiased Jaccard with variance J(1-J)/k; LSH banding; containment is a different quantity\n",
+    "`Jhat = (1/k) sum_i 1{minhash_i(A) = minhash_i(B)}`.  For a uniformly random permutation `P[min equal] = J`, so `k Jhat ~ Bin(k, J)`: unbiased, `Var = J(1-J)/k`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "source": [
+    "def jaccard(A, B): return len(A & B) / len(A | B)\n",
+    "U = 3000; rng = np.random.default_rng(7)\n",
+    "A = set(rng.choice(U, 400, replace=False).tolist()); share = 240\n",
+    "B = set(list(A)[:share]) | set(rng.choice([u for u in range(U) if u not in A], 160, replace=False).tolist())\n",
+    "J = jaccard(A, B)\n",
+    "def minhash_perm(S, perms): return np.array([min(perm[x] for x in S) for perm in perms])\n",
+    "k = 128; reps = 3000\n",
+    "perm_bank = np.array([rng.permutation(U) for _ in range(k)])           # one fresh set of k permutations per replication below\n",
+    "ests = []\n",
+    "Al = np.array(sorted(A)); Bl = np.array(sorted(B))\n",
+    "for r in range(reps):\n",
+    "    perms = rng.permuted(np.tile(np.arange(U), (k, 1)), axis=1)\n",
+    "    ests.append((perms[:, Al].min(1) == perms[:, Bl].min(1)).mean())\n",
+    "ests = np.array(ests)\n",
+    "check(\"4. E[Jhat] = J (random permutations, k=128, 3000 replications)\", abs(ests.mean() - J), 4 * math.sqrt(J * (1 - J) / k / reps), f\"J={J:.4f}\")\n",
+    "check(\"4. Var(Jhat) = J(1-J)/k\", abs(ests.var(ddof=1) / (J * (1 - J) / k) - 1), 4 * math.sqrt(2 / (reps - 1)), f\"emp {ests.var(ddof=1):.3e} vs {J*(1-J)/k:.3e}\")\n",
+    "# practical universal hash h(x) = (a x + b) mod p  (not truly random permutations)\n",
+    "P = 2_147_483_647\n",
+    "def uh_est(A_, B_, k, rg):\n",
+    "    a = rg.integers(1, P, k); b = rg.integers(0, P, k)\n",
+    "    Aa = np.array(sorted(A_), dtype=np.int64); Bb = np.array(sorted(B_), dtype=np.int64)\n",
+    "    ha = ((a[:, None] * Aa[None, :] + b[:, None]) % P).min(1); hb = ((a[:, None] * Bb[None, :] + b[:, None]) % P).min(1)\n",
+    "    return (ha == hb).mean()\n",
+    "rg = np.random.default_rng(11)\n",
+    "est_uh = np.array([uh_est(A, B, 128, rg) for _ in range(3000)])\n",
+    "# a linear hash is pairwise independent but NOT min-wise independent: a small bias is allowed and reported, not hidden\n",
+    "check(\"4. universal-hash MinHash (practical): |bias| < 0.01 (small but need not be zero; reported)\", abs(est_uh.mean() - J), 0.01, f\"bias {est_uh.mean()-J:+.4f} (MC s.e. {math.sqrt(J*(1-J)/128/3000):.4f})\")\n",
+    "for eps in (0.05, 0.02, 0.01):\n",
+    "    kreq = math.ceil(stats.norm.ppf(0.975) ** 2 * 0.25 / eps ** 2); print(f\"  worst-case (J=1/2) hashes for +-{eps} at 95%: k >= {kreq}\")\n",
+    "kk = 385; cover = 0.0\n",
+    "pm = stats.binom.pmf(np.arange(kk + 1), kk, 0.5); jh = np.arange(kk + 1) / kk\n",
+    "cover = pm[np.abs(jh - 0.5) <= 0.05].sum()\n",
+    "check(\"4. k=385 gives >= ~95% coverage of +-0.05 at J=1/2 (exact binomial)\", 0.0 if cover > 0.94 else 1.0, 0.0, f\"coverage {cover:.4f}\")\n",
+    "# decision at a decontamination threshold: flag if Jhat >= 0.8, k = 128\n",
+    "kk = 128\n",
+    "for Jt in (0.7, 0.75, 0.8, 0.85):\n",
+    "    pflag = stats.binom.sf(math.ceil(0.8 * kk) - 1, kk, Jt)\n",
+    "    print(f\"  true J={Jt}: P[flag at Jhat>=0.8] = {pflag:.3f}   (sd of Jhat {math.sqrt(Jt*(1-Jt)/kk):.3f})\")\n",
+    "Jt = 0.75; mc = (rng.binomial(kk, Jt, 2_000_000) / kk >= 0.8).mean(); ex = stats.binom.sf(math.ceil(0.8 * kk) - 1, kk, Jt)\n",
+    "check(\"4. exact flag probability at J=0.75 vs Monte Carlo\", abs(mc - ex), 4 * math.sqrt(ex * (1 - ex) / 2e6), f\"P[flag]={ex:.3f}: a 'near-duplicate' just under the threshold is flagged often\")\n",
+    "# LSH banding: b bands of r rows: P[candidate] = 1 - (1 - J^r)^b\n",
+    "b_, r_ = 20, 5\n",
+    "for Jt in (0.3, 0.5, 0.6, 0.8):\n",
+    "    th = 1 - (1 - Jt ** r_) ** b_\n",
+    "    bits = rng.random((400_000, b_, r_)) < Jt       # each hash agrees w.p. J independently (random permutations)\n",
+    "    mc = bits.all(2).any(1).mean()\n",
+    "    check(f\"4. LSH banding P[candidate] at J={Jt}: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5)\", abs(mc - th), 4 * math.sqrt(th * (1 - th) / 400_000) + 1e-4, f\"{th:.4f}\")\n",
+    "print(f\"  LSH threshold approx (1/b)^(1/r) = {(1/b_)**(1/r_):.3f}\")\n",
+    "# containment: a benchmark item (40 shingles) inside a long web page (2000 shingles): J tiny, containment 1\n",
+    "Aitem = set(range(40)); Bpage = set(range(2000))\n",
+    "Jc = jaccard(Aitem, Bpage); cont = len(Aitem & Bpage) / len(Aitem)\n",
+    "inter_from_J = Jc * (len(Aitem) + len(Bpage)) / (1 + Jc)\n",
+    "check(\"4. |A n B| = J(|A|+|B|)/(1+J): containment recoverable from Jaccard and the two sizes\", abs(inter_from_J - 40), 1e-9, f\"J={Jc:.4f} (flag threshold 0.8 would miss it), containment={cont:.2f}\")\n",
+    "# error amplification of the containment estimate via Jhat\n",
+    "kk = 128; Jh = rng.binomial(kk, Jc, 100_000) / kk\n",
+    "cont_hat = Jh * (40 + 2000) / (1 + Jh) / 40\n",
+    "print(f\"  containment estimate from MinHash at k=128: mean {cont_hat.mean():.3f}, sd {cont_hat.std():.3f} (true 1.0): very noisy for tiny J\")\n",
+    "check(\"4. containment from MinHash at k=128 for tiny J has relative sd > 0.5\", 0.0 if cont_hat.std() > 0.5 else 1.0, 0.0, f\"sd {cont_hat.std():.2f}\")\n",
+    "# real text shingles: document with edits\n",
+    "rg = np.random.default_rng(5)\n",
+    "vocab = [f\"w{i}\" for i in range(3000)]\n",
+    "doc = list(rg.choice(vocab, 600)); doc2 = doc.copy()\n",
+    "for pos in rg.choice(600, 36, replace=False): doc2[pos] = rg.choice(vocab)\n",
+    "sh = lambda d, w=5: {\" \".join(d[i:i + w]) for i in range(len(d) - w + 1)}\n",
+    "S1, S2 = sh(doc), sh(doc2); Jx = jaccard(S1, S2)\n",
+    "hs = lambda S: {int.from_bytes(__import__('hashlib').blake2b(x.encode(), digest_size=8).digest(), 'little') % P for x in S}\n",
+    "H1, H2 = hs(S1), hs(S2)\n",
+    "est = np.array([uh_est(H1, H2, 256, rg) for _ in range(400)])\n",
+    "check(\"4. 5-word shingles of a 600-word doc with 6% word edits: MinHash(k=256) mean vs exact Jaccard\", abs(est.mean() - Jx), 4 * math.sqrt(Jx * (1 - Jx) / 256 / 400) + 0.003, f\"exact J={Jx:.4f}, est {est.mean():.4f}, sd {est.std():.4f}\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "[ok] 4. E[Jhat] = J (random permutations, k=128, 3000 replications)                                           err=4.416e-04 (tol 3.2e-03) J=0.4286\n",
+      "[ok] 4. Var(Jhat) = J(1-J)/k                                                                                  err=3.114e-02 (tol 1.0e-01) emp 1.973e-03 vs 1.913e-03\n",
+      "[ok] 4. universal-hash MinHash (practical): |bias| < 0.01 (small but need not be zero; reported)              err=3.666e-03 (tol 1.0e-02) bias +0.0037 (MC s.e. 0.0008)\n",
+      "  worst-case (J=1/2) hashes for +-0.05 at 95%: k >= 385\n",
+      "  worst-case (J=1/2) hashes for +-0.02 at 95%: k >= 2401\n",
+      "  worst-case (J=1/2) hashes for +-0.01 at 95%: k >= 9604\n",
+      "[ok] 4. k=385 gives >= ~95% coverage of +-0.05 at J=1/2 (exact binomial)                                      err=0.000e+00 (tol 0.0e+00) coverage 0.9474\n",
+      "  true J=0.7: P[flag at Jhat>=0.8] = 0.005   (sd of Jhat 0.041)\n",
+      "  true J=0.75: P[flag at Jhat>=0.8] = 0.090   (sd of Jhat 0.038)\n",
+      "  true J=0.8: P[flag at Jhat>=0.8] = 0.500   (sd of Jhat 0.035)\n",
+      "  true J=0.85: P[flag at Jhat>=0.8] = 0.936   (sd of Jhat 0.032)\n",
+      "[ok] 4. exact flag probability at J=0.75 vs Monte Carlo                                                       err=1.523e-04 (tol 8.1e-04) P[flag]=0.090: a 'near-duplicate' just under the threshold is flagged often\n",
+      "[ok] 4. LSH banding P[candidate] at J=0.3: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5)                             err=2.643e-04 (tol 1.4e-03) 0.0475\n",
+      "[ok] 4. LSH banding P[candidate] at J=0.5: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5)                             err=2.007e-04 (tol 3.3e-03) 0.4701\n",
+      "[ok] 4. LSH banding P[candidate] at J=0.6: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5)                             err=2.825e-04 (tol 2.6e-03) 0.8019\n",
+      "[ok] 4. LSH banding P[candidate] at J=0.8: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5)                             err=2.394e-05 (tol 2.2e-04) 0.9996\n",
+      "  LSH threshold approx (1/b)^(1/r) = 0.549\n",
+      "[ok] 4. |A n B| = J(|A|+|B|)/(1+J): containment recoverable from Jaccard and the two sizes                    err=0.000e+00 (tol 1.0e-09) J=0.0200 (flag threshold 0.8 would miss it), containment=1.00\n",
+      "  containment estimate from MinHash at k=128: mean 0.997, sd 0.603 (true 1.0): very noisy for tiny J\n",
+      "[ok] 4. containment from MinHash at k=128 for tiny J has relative sd > 0.5                                    err=0.000e+00 (tol 0.0e+00) sd 0.60\n",
+      "[ok] 4. 5-word shingles of a 600-word doc with 6% word edits: MinHash(k=256) mean vs exact Jaccard            err=2.862e-03 (tol 9.2e-03) exact J=0.5872, est 0.5844, sd 0.0322\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Membership inference from NLL gaps\n",
+    "Members have lower NLL: `NLL_member ~ N(mu - Delta, sigma^2)`, non-member `N(mu, sigma^2)`.  AUC of the threshold test = `Phi(Delta / (sigma sqrt 2))`; TPR at FPR `f`: `Phi(Delta/sigma - z_{1-f})`.\n",
+    "If examples differ in intrinsic difficulty `D_i ~ N(0, sD^2)` (added to every NLL), the raw-NLL AUC drops to `Phi(Delta/sqrt(2(sD^2+sigma^2)))`; a reference-calibrated score `NLL_target - NLL_ref` cancels `D_i`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "source": [
+    "rng = np.random.default_rng(21); R_ = 400_000\n",
+    "Delta, sig, sD, sref = 0.30, 0.25, 0.80, 0.25\n",
+    "mem = rng.normal(-Delta, sig, R_); non = rng.normal(0, sig, R_)\n",
+    "def auc_of(pos_low, neg):\n",
+    "    # P[score(member) < score(non-member)] with MC estimate by ranks on a subsample\n",
+    "    x = pos_low[:200000]; y = neg[:200000]\n",
+    "    r = stats.rankdata(np.concatenate([x, y])); return 1 - (r[:len(x)].sum() - len(x) * (len(x) + 1) / 2) / (len(x) * len(y))\n",
+    "check(\"5. AUC = Phi(Delta/(sigma sqrt2)) (Gaussian NLL gap)\", abs(auc_of(mem, non) - stats.norm.cdf(Delta / (sig * math.sqrt(2)))), 0.003, f\"{stats.norm.cdf(Delta/(sig*math.sqrt(2))):.4f}\")\n",
+    "thr = np.quantile(non, 0.01); tpr_mc = (mem <= thr).mean(); tpr_th = stats.norm.cdf(Delta / sig - stats.norm.ppf(0.99))\n",
+    "check(\"5. TPR at 1% FPR = Phi(Delta/sigma - z_{0.99})\", abs(tpr_mc - tpr_th), 0.004, f\"{tpr_th:.4f}\")\n",
+    "Dm = rng.normal(0, sD, R_); Dn = rng.normal(0, sD, R_)\n",
+    "raw_m = Dm + mem; raw_n = Dn + non\n",
+    "auc_raw_th = stats.norm.cdf(Delta / math.sqrt(2 * (sD ** 2 + sig ** 2)))\n",
+    "check(\"5. difficulty heterogeneity: raw-NLL AUC = Phi(Delta/sqrt(2(sD^2+sigma^2)))\", abs(auc_of(raw_m, raw_n) - auc_raw_th), 0.003, f\"{auc_raw_th:.4f}\")\n",
+    "ref_m = Dm + rng.normal(0, sref, R_); ref_n = Dn + rng.normal(0, sref, R_)\n",
+    "cal_th = stats.norm.cdf(Delta / math.sqrt(2 * (sig ** 2 + sref ** 2)))\n",
+    "check(\"5. calibrated NLL - NLL_ref: AUC = Phi(Delta/sqrt(2(sigma^2+sref^2)))\", abs(auc_of(raw_m - ref_m, raw_n - ref_n) - cal_th), 0.003, f\"{cal_th:.4f}\")\n",
+    "print(f\"  AUC: raw {auc_raw_th:.3f} -> calibrated {cal_th:.3f}  (Delta={Delta}, sigma={sig}, sD={sD}, sref={sref})\")\n",
+    "# a real (tiny) language model: character n-gram, memorisation grows with order; ground-truth membership known\n",
+    "V = 12; L = 120; ndoc = 2000; alpha = 0.1; rng = np.random.default_rng(33)\n",
+    "logits = rng.normal(0, 1.5, (V, V, V))                                   # order-2 Markov source\n",
+    "def gen_docs(nd, rg):\n",
+    "    taus = np.exp(rg.uniform(math.log(0.4), math.log(2.5), nd)); X = np.zeros((nd, L), dtype=np.int64); X[:, :2] = rg.integers(0, V, (nd, 2))\n",
+    "    for t in range(2, L):\n",
+    "        lg = logits[X[:, t - 2], X[:, t - 1]] / taus[:, None]; pr = np.exp(lg - lg.max(1, keepdims=True)); pr /= pr.sum(1, keepdims=True)\n",
+    "        u = rg.random(nd)[:, None]; X[:, t] = (pr.cumsum(1) < u).sum(1).clip(max=V - 1)\n",
+    "    return X, taus\n",
+    "Xm, tm = gen_docs(ndoc, rng); Xn, tn = gen_docs(ndoc, rng); Xr, _ = gen_docs(ndoc, rng)\n",
+    "def fit(X, order):\n",
+    "    cnt = np.zeros((V ** order, V))\n",
+    "    for t in range(4, L):\n",
+    "        ctx = np.zeros(len(X), dtype=np.int64)\n",
+    "        for j in range(order, 0, -1): ctx = ctx * V + X[:, t - j]\n",
+    "        np.add.at(cnt, (ctx, X[:, t]), 1)\n",
+    "    return cnt\n",
+    "def nll(cnt, X, order):\n",
+    "    out = np.zeros(len(X))\n",
+    "    for t in range(4, L):\n",
+    "        ctx = np.zeros(len(X), dtype=np.int64)\n",
+    "        for j in range(order, 0, -1): ctx = ctx * V + X[:, t - j]\n",
+    "        c = cnt[ctx]; out -= np.log((c[np.arange(len(X)), X[:, t]] + alpha) / (c.sum(1) + alpha * V))\n",
+    "    return out / (L - 4)\n",
+    "def auc_small(low, high):  # P[low < high] (+0.5 ties)\n",
+    "    r = stats.rankdata(np.concatenate([low, high])); n1 = len(low); return 1 - (r[:n1].sum() - n1 * (n1 + 1) / 2) / (n1 * len(high))\n",
+    "res = {}\n",
+    "for order in (1, 2, 3, 4):\n",
+    "    cm = fit(Xm, order); cr = fit(Xr, order)\n",
+    "    nm, nn = nll(cm, Xm, order), nll(cm, Xn, order)\n",
+    "    rm, rn = nll(cr, Xm, order), nll(cr, Xn, order)\n",
+    "    res[order] = (auc_small(nm, nn), auc_small(nm - rm, nn - rn), nn.mean() - nm.mean())\n",
+    "    print(f\"  order {order}: raw-NLL AUC = {res[order][0]:.3f}, calibrated AUC = {res[order][1]:.3f}, mean NLL gap (non - member) = {res[order][2]:.4f} nats/char\")\n",
+    "check(\"5. no memorisation capacity at order 1-2 (<= true order): raw AUC within 0.04 of 0.5\", max(abs(res[1][0] - 0.5), abs(res[2][0] - 0.5)), 0.04)\n",
+    "check(\"5. membership signal grows with model capacity: raw AUC order 4 > order 3 > order 2\", 0.0 if res[4][0] > res[3][0] > res[2][0] else 1.0, 0.0)\n",
+    "check(\"5. reference calibration helps at the high-capacity order: calibrated AUC > raw AUC + 0.05 (order 4)\", 0.0 if res[4][1] > res[4][0] + 0.05 else 1.0, 0.0, f\"{res[4][0]:.3f} -> {res[4][1]:.3f}\")\n",
+    "check(\"5. calibrated AUC increases with order and exceeds raw AUC at every order >= 2 (it detects even the small order-2 overfit that raw NLL hides)\", 0.0 if (res[1][1] < res[2][1] < res[3][1] < res[4][1] and all(res[o][1] > res[o][0] for o in (2, 3, 4))) else 1.0, 0.0)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "[ok] 5. AUC = Phi(Delta/(sigma sqrt2)) (Gaussian NLL gap)                                                     err=3.136e-05 (tol 3.0e-03) 0.8019\n",
+      "[ok] 5. TPR at 1% FPR = Phi(Delta/sigma - z_{0.99})                                                           err=1.543e-03 (tol 4.0e-03) 0.1300\n",
+      "[ok] 5. difficulty heterogeneity: raw-NLL AUC = Phi(Delta/sqrt(2(sD^2+sigma^2)))                              err=1.669e-04 (tol 3.0e-03) 0.5999\n",
+      "[ok] 5. calibrated NLL - NLL_ref: AUC = Phi(Delta/sqrt(2(sigma^2+sref^2)))                                    err=7.603e-04 (tol 3.0e-03) 0.7257\n",
+      "  AUC: raw 0.600 -> calibrated 0.726  (Delta=0.3, sigma=0.25, sD=0.8, sref=0.25)\n",
+      "  order 1: raw-NLL AUC = 0.506, calibrated AUC = 0.555, mean NLL gap (non - member) = 0.0019 nats/char\n",
+      "  order 2: raw-NLL AUC = 0.510, calibrated AUC = 0.675, mean NLL gap (non - member) = 0.0154 nats/char\n",
+      "  order 3: raw-NLL AUC = 0.559, calibrated AUC = 0.832, mean NLL gap (non - member) = 0.1048 nats/char\n",
+      "  order 4: raw-NLL AUC = 0.715, calibrated AUC = 0.923, mean NLL gap (non - member) = 0.4656 nats/char\n",
+      "[ok] 5. no memorisation capacity at order 1-2 (<= true order): raw AUC within 0.04 of 0.5                     err=1.007e-02 (tol 4.0e-02) \n",
+      "[ok] 5. membership signal grows with model capacity: raw AUC order 4 > order 3 > order 2                      err=0.000e+00 (tol 0.0e+00) \n",
+      "[ok] 5. reference calibration helps at the high-capacity order: calibrated AUC > raw AUC + 0.05 (order 4)     err=0.000e+00 (tol 0.0e+00) 0.715 -> 0.923\n",
+      "[ok] 5. calibrated AUC increases with order and exceeds raw AUC at every order >= 2 (it detects even the small order-2 overfit that raw NLL hides) err=0.000e+00 (tol 0.0e+00) \n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Summary"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "source": [
+    "for n_, e_, t_ in CHECKS: print(f\"PASS {n_[:110]}\")\n",
+    "print(f\"\\n{len(CHECKS)} checks, all passed\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "PASS 1. precision closed form vs Monte Carlo (4e6 pages)\n",
+      "PASS 1. recall closed form vs Monte Carlo\n",
+      "PASS 1. same classifier, same threshold: precision at pi=0.2 exceeds precision at pi=0.01 by > 0.4\n",
+      "PASS 1. argmax of the net-gain curve vs closed-form posterior-cutoff threshold\n",
+      "PASS 1. posterior P(rel | s=t*) = h/(g+h)\n",
+      "PASS 1. the F1-optimal threshold is not the gain-optimal one (differ by > 0.2) and loses downstream gain\n",
+      "PASS 1. iterative re-seeding: recall rises monotonically and increments shrink geometrically with ratio (1-gamma)\n",
+      "PASS 2. E[ghat] = (1-1/n) p(1-p) exactly (7 values of p)\n",
+      "PASS 2. same-sample selection weights exceed independent-pilot weights for every p in the table (selection on the s\n",
+      "PASS 2. w_same(p=0.1) exact vs 2e6 Monte Carlo\n",
+      "PASS 2. beta-binomial exact P[band] vs MC\n",
+      "PASS 2. a substantial share of selected problems lie outside the band in TRUE pass rate (> 15%)\n",
+      "PASS 3. E[fresh pass rate | first-pass count c=0] = (a+c)/(a+b+n)\n",
+      "PASS 3. E[fresh pass rate | first-pass count c=1] = (a+c)/(a+b+n)\n",
+      "PASS 3. E[fresh pass rate | first-pass count c=2] = (a+c)/(a+b+n)\n",
+      "PASS 3. naive before/after overstates the effect (regression to the mean) and DiD recovers it\n",
+      "PASS 3. naive before/after is inflated by RTM: excess over truth exceeds 0.03\n",
+      "PASS 4. E[Jhat] = J (random permutations, k=128, 3000 replications)\n",
+      "PASS 4. Var(Jhat) = J(1-J)/k\n",
+      "PASS 4. universal-hash MinHash (practical): |bias| < 0.01 (small but need not be zero; reported)\n",
+      "PASS 4. k=385 gives >= ~95% coverage of +-0.05 at J=1/2 (exact binomial)\n",
+      "PASS 4. exact flag probability at J=0.75 vs Monte Carlo\n",
+      "PASS 4. LSH banding P[candidate] at J=0.3: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5)\n",
+      "PASS 4. LSH banding P[candidate] at J=0.5: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5)\n",
+      "PASS 4. LSH banding P[candidate] at J=0.6: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5)\n",
+      "PASS 4. LSH banding P[candidate] at J=0.8: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5)\n",
+      "PASS 4. |A n B| = J(|A|+|B|)/(1+J): containment recoverable from Jaccard and the two sizes\n",
+      "PASS 4. containment from MinHash at k=128 for tiny J has relative sd > 0.5\n",
+      "PASS 4. 5-word shingles of a 600-word doc with 6% word edits: MinHash(k=256) mean vs exact Jaccard\n",
+      "PASS 5. AUC = Phi(Delta/(sigma sqrt2)) (Gaussian NLL gap)\n",
+      "PASS 5. TPR at 1% FPR = Phi(Delta/sigma - z_{0.99})\n",
+      "PASS 5. difficulty heterogeneity: raw-NLL AUC = Phi(Delta/sqrt(2(sD^2+sigma^2)))\n",
+      "PASS 5. calibrated NLL - NLL_ref: AUC = Phi(Delta/sqrt(2(sigma^2+sref^2)))\n",
+      "PASS 5. no memorisation capacity at order 1-2 (<= true order): raw AUC within 0.04 of 0.5\n",
+      "PASS 5. membership signal grows with model capacity: raw AUC order 4 > order 3 > order 2\n",
+      "PASS 5. reference calibration helps at the high-capacity order: calibrated AUC > raw AUC + 0.05 (order 4)\n",
+      "PASS 5. calibrated AUC increases with order and exceeds raw AUC at every order >= 2 (it detects even the small orde\n",
+      "\n",
+      "37 checks, all passed\n"
+     ]
+    }
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

code/notebooks/0301-curry-howard.ipynb

@@ -0,0 +1,773 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0999126c",
+   "metadata": {},
+   "source": [
+    "# 0301 - Curry-Howard: propositions as types, proofs as terms\n",
+    "\n",
+    "Scratch-pad for note `0301 - Curry-Howard.pdf`. Same discipline as the Flash-Attention-101 scratch-pads:\n",
+    "the **naive** implementation first, then the **better** one, then an assertion that they agree.\n",
+    "\n",
+    "Plan\n",
+    "1. A 40-line kernel for the simply typed lambda calculus with `->` and `False` (`infer`, `normalize`).\n",
+    "2. Subject reduction (Boxed result 2) checked on thousands of random typed terms.\n",
+    "3. Inhabitation, NAIVE: depth-bounded enumeration of beta-normal terms. EFFICIENT: Dyckhoff's contraction-free\n",
+    "   calculus G4ip with proof-term extraction. Assert the verdicts agree on every formula with <= 4 arrows.\n",
+    "4. Soundness (provable => classically valid, provable => forced in every Kripke chain) and the gap\n",
+    "   (Peirce, double-negation elimination: classically valid, not inhabited).\n",
+    "5. What the Lean kernel actually sees: `#print axioms` for each way of \"proving\" something; the axiom gate.\n",
+    "\n",
+    "No GPU is used (`CUDA_VISIBLE_DEVICES` is cleared). Lean calls go through `align.verify.lean` with ONE worker, nice 15."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "d045e611",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:53:55.822833Z",
+     "iopub.status.busy": "2026-10-14T09:53:55.822602Z",
+     "iopub.status.idle": "2026-10-14T09:53:55.843322Z",
+     "shell.execute_reply": "2026-10-14T09:53:55.842326Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "K : A → B → A\n",
+      "S : (A → B → C) → (A → B) → A → C\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
+    "import itertools, random, time, math, sys\n",
+    "from collections import Counter\n",
+    "random.seed(0)\n",
+    "\n",
+    "# Types:  ('atom', name) | ('bot',) | ('arr', A, B)        Terms: ('var',x) | ('lam',x,A,t) | ('app',f,u) | ('abort',A,t)\n",
+    "BOT = ('bot',)\n",
+    "def atom(n): return ('atom', n)\n",
+    "def arr(a, b): return ('arr', a, b)\n",
+    "def show_ty(t):\n",
+    "    if t[0] == 'atom': return t[1]\n",
+    "    if t[0] == 'bot': return '⊥'\n",
+    "    l = show_ty(t[1]); l = f'({l})' if t[1][0] == 'arr' else l\n",
+    "    return f'{l} → {show_ty(t[2])}'\n",
+    "def show_tm(t):\n",
+    "    if t[0] == 'var': return t[1]\n",
+    "    if t[0] == 'lam': return f'(λ{t[1]}:{show_ty(t[2])}. {show_tm(t[3])})'\n",
+    "    if t[0] == 'app': return f'({show_tm(t[1])} {show_tm(t[2])})'\n",
+    "    return f'(abort[{show_ty(t[1])}] {show_tm(t[2])})'\n",
+    "\n",
+    "class TypeErr(Exception): pass\n",
+    "\n",
+    "def infer(ctx, t):\n",
+    "    \"\"\"Typing judgement  ctx |- t : A  (ctx: dict var->type). Rules: Var, ->I, ->E, bot-E.\"\"\"\n",
+    "    k = t[0]\n",
+    "    if k == 'var':\n",
+    "        if t[1] not in ctx: raise TypeErr(f'unbound {t[1]}')\n",
+    "        return ctx[t[1]]\n",
+    "    if k == 'lam':\n",
+    "        _, x, A, body = t\n",
+    "        return arr(A, infer({**ctx, x: A}, body))\n",
+    "    if k == 'app':\n",
+    "        F, U = infer(ctx, t[1]), infer(ctx, t[2])\n",
+    "        if F[0] != 'arr' or F[1] != U: raise TypeErr(f'cannot apply {show_ty(F)} to {show_ty(U)}')\n",
+    "        return F[2]\n",
+    "    if k == 'abort':\n",
+    "        if infer(ctx, t[2]) != BOT: raise TypeErr('abort needs a proof of ⊥')\n",
+    "        return t[1]\n",
+    "    raise TypeErr(k)\n",
+    "\n",
+    "_fresh = itertools.count()\n",
+    "KINDS = Counter()                                  # which reduction rule fired (coverage of the SR check)\n",
+    "def fv(t):\n",
+    "    k = t[0]\n",
+    "    if k == 'var': return {t[1]}\n",
+    "    if k == 'lam': return fv(t[3]) - {t[1]}\n",
+    "    if k == 'app': return fv(t[1]) | fv(t[2])\n",
+    "    return fv(t[2])\n",
+    "\n",
+    "def subst(t, x, s):\n",
+    "    \"\"\"Capture-avoiding substitution t[x := s].\"\"\"\n",
+    "    k = t[0]\n",
+    "    if k == 'var': return s if t[1] == x else t\n",
+    "    if k == 'app': return ('app', subst(t[1], x, s), subst(t[2], x, s))\n",
+    "    if k == 'abort': return ('abort', t[1], subst(t[2], x, s))\n",
+    "    _, y, A, b = t\n",
+    "    if y == x: return t\n",
+    "    if y in fv(s):\n",
+    "        z = f'{y}_{next(_fresh)}'\n",
+    "        b = subst(b, y, ('var', z)); y = z\n",
+    "    return ('lam', y, A, subst(b, x, s))\n",
+    "\n",
+    "def step(t):\n",
+    "    \"\"\"One leftmost-outermost step: beta, and the commuting conversion (abort[A->B] e) u ~> abort[B] e. None if normal.\"\"\"\n",
+    "    k = t[0]\n",
+    "    if k == 'app':\n",
+    "        f, u = t[1], t[2]\n",
+    "        if f[0] == 'lam': KINDS['beta'] += 1; return subst(f[3], f[1], u)\n",
+    "        if f[0] == 'abort' and f[1][0] == 'arr': KINDS['comm'] += 1; return ('abort', f[1][2], f[2])\n",
+    "        r = step(f)\n",
+    "        if r is not None: return ('app', r, u)\n",
+    "        r = step(u)\n",
+    "        return None if r is None else ('app', f, r)\n",
+    "    if k == 'lam':\n",
+    "        r = step(t[3]); return None if r is None else ('lam', t[1], t[2], r)\n",
+    "    if k == 'abort':\n",
+    "        r = step(t[2]); return None if r is None else ('abort', t[1], r)\n",
+    "    return None\n",
+    "\n",
+    "def normalize(t, fuel=10_000):\n",
+    "    for _ in range(fuel):\n",
+    "        r = step(t)\n",
+    "        if r is None: return t\n",
+    "        t = r\n",
+    "    raise RuntimeError('fuel exhausted (would contradict strong normalisation)')\n",
+    "\n",
+    "# the classic combinators, as terms, with their types read off by the kernel\n",
+    "A, B, C = atom('A'), atom('B'), atom('C')\n",
+    "K = ('lam','x',A,('lam','y',B,('var','x')))\n",
+    "S = ('lam','f',arr(A,arr(B,C)),('lam','g',arr(A,B),('lam','x',A,('app',('app',('var','f'),('var','x')),('app',('var','g'),('var','x'))))))\n",
+    "for name, tm in [('K', K), ('S', S)]:\n",
+    "    print(name, ':', show_ty(infer({}, tm)))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "057302f6",
+   "metadata": {},
+   "source": [
+    "## 2. Subject reduction (Boxed result 2): if `Gamma |- t : A` and `t ~> t'` then `Gamma |- t' : A`\n",
+    "Check on random well-typed terms: generate by random type-directed construction, reduce step by step, re-infer at every step."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "5787fcb2",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:53:55.845545Z",
+     "iopub.status.busy": "2026-10-14T09:53:55.845323Z",
+     "iopub.status.idle": "2026-10-14T09:53:56.359730Z",
+     "shell.execute_reply": "2026-10-14T09:53:56.358821Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "3000 random typed terms, 1205 with >=1 redex, 1851 reduction steps (1431 beta, 420 commuting-conversion); subject reduction violations: 0  (0.5s)\n"
+     ]
+    }
+   ],
+   "source": [
+    "ATOMS = ['a', 'b', 'c']\n",
+    "def rand_type(d):\n",
+    "    if d == 0 or random.random() < 0.35:\n",
+    "        return random.choice([atom(x) for x in ATOMS] + [BOT])\n",
+    "    return arr(rand_type(d - 1), rand_type(d - 1))\n",
+    "\n",
+    "def spine_end(T):\n",
+    "    while T[0] == 'arr': T = T[2]\n",
+    "    return T\n",
+    "\n",
+    "def rand_term(ctx, goal, d):\n",
+    "    \"\"\"random term of type `goal` in `ctx`, possibly with redexes (introduces cuts).\"\"\"\n",
+    "    if d > 0 and random.random() < 0.35:                      # a cut: (λx:T. body) arg\n",
+    "        T = rand_type(2)\n",
+    "        arg = rand_term(ctx, T, d - 1)\n",
+    "        x = f'v{next(_fresh)}'\n",
+    "        body = rand_term({**ctx, x: T}, goal, d - 1)\n",
+    "        return None if (arg is None or body is None) else ('app', ('lam', x, T, body), arg)\n",
+    "    if d > 0 and random.random() < 0.15 and any(spine_end(T) == BOT for T in ctx.values()):   # ex falso at an ARROW type, then applied\n",
+    "        T = rand_type(2)\n",
+    "        e = rand_term(ctx, BOT, d - 1); arg = rand_term(ctx, T, d - 1)\n",
+    "        if e is not None and arg is not None: return ('app', ('abort', arr(T, goal), e), arg)\n",
+    "    if goal[0] == 'arr':\n",
+    "        x = f'v{next(_fresh)}'\n",
+    "        body = rand_term({**ctx, x: goal[1]}, goal[2], d - 1)\n",
+    "        return None if body is None else ('lam', x, goal[1], body)\n",
+    "    # atomic goal: use a hypothesis head with matching result, else give up\n",
+    "    cands = []\n",
+    "    for x, T in ctx.items():\n",
+    "        spine, R = [], T\n",
+    "        while R[0] == 'arr': spine.append(R[1]); R = R[2]\n",
+    "        if R == goal or (R == BOT):\n",
+    "            cands.append((x, spine, R))\n",
+    "    if not cands: return None\n",
+    "    x, spine, R = random.choice(cands)\n",
+    "    args = []\n",
+    "    for Ai in spine:\n",
+    "        a = rand_term(ctx, Ai, d - 1) if d > 0 else None\n",
+    "        if a is None: return None\n",
+    "        args.append(a)\n",
+    "    tm = ('var', x)\n",
+    "    for a in args: tm = ('app', tm, a)\n",
+    "    return tm if R == goal else ('abort', goal, tm)\n",
+    "\n",
+    "n_ok = n_steps = n_nontrivial = 0\n",
+    "KINDS.clear()\n",
+    "t0 = time.time()\n",
+    "while n_ok < 3000:\n",
+    "    ctx = {f'h{i}': rand_type(2) for i in range(random.randint(0, 3))}\n",
+    "    goal = rand_type(2)\n",
+    "    tm = rand_term(ctx, goal, 4)\n",
+    "    if tm is None: continue\n",
+    "    ty = infer(ctx, tm)                       # must typecheck by construction\n",
+    "    assert ty == goal\n",
+    "    cur, steps = tm, 0\n",
+    "    while True:\n",
+    "        nxt = step(cur)\n",
+    "        if nxt is None: break\n",
+    "        assert infer(ctx, nxt) == goal, ('subject reduction violated', show_tm(cur), show_tm(nxt))\n",
+    "        cur, steps = nxt, steps + 1\n",
+    "        n_steps += 1\n",
+    "    n_nontrivial += steps > 0\n",
+    "    n_ok += 1\n",
+    "print(f'{n_ok} random typed terms, {n_nontrivial} with >=1 redex, {n_steps} reduction steps '\n",
+    "      f'({KINDS[\"beta\"]} beta, {KINDS[\"comm\"]} commuting-conversion); subject reduction violations: 0  ({time.time()-t0:.1f}s)')\n",
+    "assert KINDS['comm'] > 0 and KINDS['beta'] > 0          # both cases of Theorem 2 are actually exercised"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fc48f5ca",
+   "metadata": {},
+   "source": [
+    "## 3. Inhabitation: naive vs efficient\n",
+    "**Naive** (`search_naive`): enumerate beta-normal, eta-long terms by depth, backtracking. Correct but exponential, and it can only answer\n",
+    "\"found within depth D\" / \"not found within depth D\".\n",
+    "\n",
+    "**Efficient** (`g4ip`): Dyckhoff's contraction-free sequent calculus (terminating, so a *negative* answer is a proof of\n",
+    "non-provability), with the proof term extracted and re-checked by the independent kernel `infer`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "b941a9af",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:53:56.362552Z",
+     "iopub.status.busy": "2026-10-14T09:53:56.362347Z",
+     "iopub.status.idle": "2026-10-14T09:53:56.376383Z",
+     "shell.execute_reply": "2026-10-14T09:53:56.375472Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "K          a→b→a                         proof term: (λy0:a. (λy1:b. y0))\n",
+      "S                                        proof term: (λy2:a → b → c. (λy3:a → b. (λy4:a. ((y2 y4) (y3 y4)))))\n",
+      "contraposition (a→b)→¬b→¬a               proof term: (λy8:a → b. (λy9:b → ⊥. (λy10:a. (abort[⊥] (y9 (y8 y10))))))\n",
+      "a → ¬¬a                                  proof term: (λy13:a. (λy14:a → ⊥. (abort[⊥] (y14 y13))))\n",
+      "¬¬¬a → ¬a                                proof term: (λy16:((a → ⊥) → ⊥) → ⊥. (λy17:a. (abort[⊥] (y16 (λc18:a → ⊥. (abort[⊥] (c18 y17)))))))\n",
+      "ex falso ⊥→a                             proof term: (λy24:⊥. (abort[a] y24))\n",
+      "Peirce ((a→b)→a)→a                       UNPROVABLE (G4ip terminates: this is a refutation, no term exists)\n",
+      "DNE ¬¬a→a                                UNPROVABLE (G4ip terminates: this is a refutation, no term exists)\n",
+      "rev. contraposition (¬b→¬a)→a→b          UNPROVABLE (G4ip terminates: this is a refutation, no term exists)\n"
+     ]
+    }
+   ],
+   "source": [
+    "def spine_of(T):\n",
+    "    sp = []\n",
+    "    while T[0] == 'arr': sp.append(T[1]); T = T[2]\n",
+    "    return sp, T\n",
+    "\n",
+    "def search_naive(ctx, goal, depth):\n",
+    "    \"\"\"Return a normal term or None. ctx: tuple of (var,type).\"\"\"\n",
+    "    if depth <= 0: return None\n",
+    "    if goal[0] == 'arr':\n",
+    "        x = f'x{len(ctx)}'\n",
+    "        body = search_naive(ctx + ((x, goal[1]),), goal[2], depth - 1)\n",
+    "        return None if body is None else ('lam', x, goal[1], body)\n",
+    "    for x, T in ctx:\n",
+    "        sp, R = spine_of(T)\n",
+    "        if R != goal and R != BOT: continue\n",
+    "        args, okk = [], True\n",
+    "        for Ai in sp:\n",
+    "            a = search_naive(ctx, Ai, depth - 1)\n",
+    "            if a is None: okk = False; break\n",
+    "            args.append(a)\n",
+    "        if not okk: continue\n",
+    "        tm = ('var', x)\n",
+    "        for a in args: tm = ('app', tm, a)\n",
+    "        return tm if R == goal else ('abort', goal, tm)\n",
+    "    return None\n",
+    "\n",
+    "_nf = itertools.count()\n",
+    "def g4ip(ctx, goal, memo_fail):\n",
+    "    \"\"\"ctx: tuple of (var,type). Returns a proof term or None. Contraction-free: terminates.\"\"\"\n",
+    "    key = (frozenset(T for _, T in ctx), goal)\n",
+    "    if key in memo_fail: return None\n",
+    "    res = _g4ip(ctx, goal, memo_fail)\n",
+    "    if res is None: memo_fail.add(key)\n",
+    "    return res\n",
+    "\n",
+    "def _g4ip(ctx, goal, mf):\n",
+    "    for x, T in ctx:                                           # L⊥\n",
+    "        if T == BOT: return ('abort', goal, ('var', x))\n",
+    "    for x, T in ctx:                                           # Ax\n",
+    "        if T == goal and T[0] == 'atom': return ('var', x)\n",
+    "    if goal[0] == 'arr':                                       # R→ (invertible: no backtracking needed)\n",
+    "        y = f'y{next(_nf)}'\n",
+    "        b = g4ip(ctx + ((y, goal[1]),), goal[2], mf)\n",
+    "        return None if b is None else ('lam', y, goal[1], b)\n",
+    "    # L0→ (invertible): p, p→B in ctx  ==>  replace p→B by B\n",
+    "    for i, (f, T) in enumerate(ctx):\n",
+    "        if T[0] == 'arr' and T[1][0] == 'atom':\n",
+    "            for x, Tx in ctx:\n",
+    "                if Tx == T[1]:\n",
+    "                    z = f'z{next(_nf)}'\n",
+    "                    rest = ctx[:i] + ctx[i+1:] + ((z, T[2]),)\n",
+    "                    t = g4ip(rest, goal, mf)\n",
+    "                    return None if t is None else subst(t, z, ('app', ('var', f), ('var', x)))\n",
+    "    # L⊥→ (invertible): drop hypotheses ⊥→B (carry no information)\n",
+    "    for i, (f, T) in enumerate(ctx):\n",
+    "        if T[0] == 'arr' and T[1] == BOT:\n",
+    "            return g4ip(ctx[:i] + ctx[i+1:], goal, mf)\n",
+    "    # L→→ (NOT invertible: backtrack over the choice of hypothesis)\n",
+    "    for i, (f, T) in enumerate(ctx):\n",
+    "        if T[0] == 'arr' and T[1][0] == 'arr':\n",
+    "            C_, D_ = T[1][1], T[1][2]; B_ = T[2]\n",
+    "            c, g, z = f'c{next(_nf)}', f'g{next(_nf)}', f'z{next(_nf)}'\n",
+    "            base = ctx[:i] + ctx[i+1:]\n",
+    "            t1 = g4ip(base + ((c, C_), (g, arr(D_, B_))), D_, mf)\n",
+    "            if t1 is None: continue\n",
+    "            t2 = g4ip(base + ((z, B_),), goal, mf)\n",
+    "            if t2 is None: continue\n",
+    "            d = f'd{next(_nf)}'\n",
+    "            u = ('lam', c, C_, subst(t1, g, ('lam', d, D_, ('app', ('var', f), ('lam', f'_{next(_nf)}', C_, ('var', d))))))\n",
+    "            return subst(t2, z, ('app', ('var', f), u))\n",
+    "    return None\n",
+    "\n",
+    "def prove(goal):\n",
+    "    return g4ip((), goal, set())\n",
+    "\n",
+    "# Gentzen-style sanity on named formulas\n",
+    "a, b, c = atom('a'), atom('b'), atom('c')\n",
+    "neg = lambda p: arr(p, BOT)\n",
+    "named = {\n",
+    " 'K          a→b→a':                    arr(a, arr(b, a)),\n",
+    " 'S':                                    arr(arr(a, arr(b, c)), arr(arr(a, b), arr(a, c))),\n",
+    " 'contraposition (a→b)→¬b→¬a':           arr(arr(a, b), arr(neg(b), neg(a))),\n",
+    " 'a → ¬¬a':                              arr(a, neg(neg(a))),\n",
+    " '¬¬¬a → ¬a':                            arr(neg(neg(neg(a))), neg(a)),\n",
+    " 'ex falso ⊥→a':                         arr(BOT, a),\n",
+    " 'Peirce ((a→b)→a)→a':                   arr(arr(arr(a, b), a), a),\n",
+    " 'DNE ¬¬a→a':                            arr(neg(neg(a)), a),\n",
+    " 'rev. contraposition (¬b→¬a)→a→b':      arr(arr(neg(b), neg(a)), arr(a, b)),\n",
+    "}\n",
+    "for name, f in named.items():\n",
+    "    tm = prove(f)\n",
+    "    if tm is None:\n",
+    "        print(f'{name:40s} UNPROVABLE (G4ip terminates: this is a refutation, no term exists)')\n",
+    "    else:\n",
+    "        assert infer({}, tm) == f                       # independent kernel check of the extracted proof\n",
+    "        nf = normalize(tm); assert infer({}, nf) == f\n",
+    "        print(f'{name:40s} proof term: {show_tm(nf)}')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "3ec7fb35",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:53:56.378933Z",
+     "iopub.status.busy": "2026-10-14T09:53:56.378753Z",
+     "iopub.status.idle": "2026-10-14T09:53:57.467575Z",
+     "shell.execute_reply": "2026-10-14T09:53:57.466497Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "formulas: {0: 3, 1: 9, 2: 54, 3: 405, 4: 3402} total 3873\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "G4ip: 0.03s   naive(depth 14): 1.06s   speed-up x41.7\n",
+      "inhabited: 1775 of 3873 | verdict disagreements naive vs G4ip: 0\n"
+     ]
+    }
+   ],
+   "source": [
+    "# All implicational formulas over {a, b, ⊥} with exactly k arrows, k = 0..4\n",
+    "LEAVES = [atom('a'), atom('b'), BOT]\n",
+    "@__import__('functools').lru_cache(None)\n",
+    "def formulas(k):\n",
+    "    if k == 0: return tuple(LEAVES)\n",
+    "    out = []\n",
+    "    for i in range(k):\n",
+    "        for l in formulas(i):\n",
+    "            for r in formulas(k - 1 - i):\n",
+    "                out.append(arr(l, r))\n",
+    "    return tuple(out)\n",
+    "ALL = [f for k in range(5) for f in formulas(k)]\n",
+    "print('formulas:', {k: len(formulas(k)) for k in range(5)}, 'total', len(ALL))\n",
+    "\n",
+    "t0 = time.time(); verdict_eff = {}\n",
+    "for f in ALL:\n",
+    "    tm = prove(f)\n",
+    "    if tm is not None:\n",
+    "        assert infer({}, tm) == f\n",
+    "        assert infer({}, normalize(tm)) == f\n",
+    "    verdict_eff[f] = tm is not None\n",
+    "t_eff = time.time() - t0\n",
+    "\n",
+    "t0 = time.time(); verdict_naive = {}; D = 14\n",
+    "for f in ALL:\n",
+    "    tm = search_naive((), f, D)\n",
+    "    if tm is not None: assert infer({}, tm) == f\n",
+    "    verdict_naive[f] = tm is not None\n",
+    "t_naive = time.time() - t0\n",
+    "\n",
+    "disagree = [f for f in ALL if verdict_eff[f] != verdict_naive[f]]\n",
+    "assert not disagree, [show_ty(f) for f in disagree[:5]]\n",
+    "print(f'G4ip: {t_eff:.2f}s   naive(depth {D}): {t_naive:.2f}s   speed-up x{t_naive/max(t_eff,1e-9):.1f}')\n",
+    "print('inhabited:', sum(verdict_eff.values()), 'of', len(ALL), '| verdict disagreements naive vs G4ip: 0')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "05df5ec0",
+   "metadata": {},
+   "source": [
+    "## 4. Soundness and the gap\n",
+    "* **Classical semantics** (truth table, `⊥ = False`): every inhabited type must be a tautology (Boxed result 3, soundness direction).\n",
+    "* **Kripke semantics** over chains of <=3 worlds: every inhabited type must be forced at *every* world of *every* model;\n",
+    "  a non-inhabited type should have a countermodel (we report how many our tiny model class catches)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "8247c153",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:53:57.469657Z",
+     "iopub.status.busy": "2026-10-14T09:53:57.469490Z",
+     "iopub.status.idle": "2026-10-14T09:53:57.982720Z",
+     "shell.execute_reply": "2026-10-14T09:53:57.981661Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "provable 1775  all classically valid and forced in all 29 chain models: OK\n",
+      "unprovable 2098: classically valid but NOT inhabited = 32;  with a chain countermodel (m<=3) = 2098\n",
+      "   valid, not inhabited: ((a → b) → a) → a\n",
+      "   valid, not inhabited: ((a → b) → ⊥) → a\n",
+      "   valid, not inhabited: ((a → ⊥) → a) → a\n",
+      "   valid, not inhabited: ((a → ⊥) → ⊥) → a\n",
+      "   valid, not inhabited: ((b → a) → b) → b\n",
+      "   valid, not inhabited: ((b → a) → ⊥) → b\n",
+      "Peirce countermodel: 2 worlds, a true only at world 1, b never  -> world 0 does not force Peirce.\n"
+     ]
+    }
+   ],
+   "source": [
+    "def classical(f, val):\n",
+    "    if f[0] == 'atom': return val[f[1]]\n",
+    "    if f[0] == 'bot': return False\n",
+    "    return (not classical(f[1], val)) or classical(f[2], val)\n",
+    "def tautology(f):\n",
+    "    return all(classical(f, dict(zip('ab', v))) for v in itertools.product([False, True], repeat=2))\n",
+    "\n",
+    "def forces(f, w, m, th):\n",
+    "    \"\"\"chain of worlds 0..m-1; atom x true at w iff w >= th[x] (monotone).\"\"\"\n",
+    "    if f[0] == 'atom': return w >= th[f[1]]\n",
+    "    if f[0] == 'bot': return False\n",
+    "    return all((not forces(f[1], v, m, th)) or forces(f[2], v, m, th) for v in range(w, m))\n",
+    "\n",
+    "models = [(m, dict(zip('ab', (ta, tb)))) for m in (1, 2, 3) for ta in range(m + 1) for tb in range(m + 1)]\n",
+    "prov = [f for f in ALL if verdict_eff[f]]\n",
+    "unprov = [f for f in ALL if not verdict_eff[f]]\n",
+    "assert all(tautology(f) for f in prov)                                              # soundness (classical)\n",
+    "assert all(forces(f, w, m, th) for f in prov for (m, th) in models for w in range(m))   # soundness (Kripke)\n",
+    "valid_unprov = [f for f in unprov if tautology(f)]\n",
+    "countermodel = [f for f in unprov if any(not forces(f, 0, m, th) for (m, th) in models)]\n",
+    "print(f'provable {len(prov)}  all classically valid and forced in all {len(models)} chain models: OK')\n",
+    "print(f'unprovable {len(unprov)}: classically valid but NOT inhabited = {len(valid_unprov)};  '\n",
+    "      f'with a chain countermodel (m<=3) = {len(countermodel)}')\n",
+    "ex = sorted(valid_unprov, key=lambda f: len(show_ty(f)))[:6]\n",
+    "for f in ex: print('   valid, not inhabited:', show_ty(f))\n",
+    "peirce = arr(arr(arr(a, b), a), a)\n",
+    "assert tautology(peirce) and not verdict_eff[peirce] and not forces(peirce, 0, 2, {'a': 1, 'b': 2})\n",
+    "print('Peirce countermodel: 2 worlds, a true only at world 1, b never  -> world 0 does not force Peirce.')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7c345d80",
+   "metadata": {},
+   "source": [
+    "## 5. What the Lean kernel sees\n",
+    "Lean elaborates tactics into terms and the **kernel** re-checks the term. The only thing that can make a \"proof\" of a false\n",
+    "statement pass is an axiom in the term's dependency closure. `#print axioms` lists that closure.\n",
+    "Exactly the same tests as the strict verifier's axiom gate (`align/verify/lean/strict.py`, rule 3)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "f9b00c3c",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:53:57.984954Z",
+     "iopub.status.busy": "2026-10-14T09:53:57.984774Z",
+     "iopub.status.idle": "2026-10-14T09:54:02.121622Z",
+     "shell.execute_reply": "2026-10-14T09:54:02.120529Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "theorem      axioms                                        gate\n",
+      "k_comb       []                                            ACCEPT\n",
+      "peirce_law   ['Classical.choice', 'Quot.sound', 'propext'] ACCEPT\n",
+      "t_omega      ['Quot.sound', 'propext']                     ACCEPT\n",
+      "t_decide     ['propext']                                   ACCEPT\n",
+      "t_native     ['Lean.ofReduceBool', 'propext']              REJECT\n",
+      "t_sorry      ['sorryAx']                                   REJECT\n",
+      "t_evil       ['evil']                                      REJECT\n",
+      "t_funext     ['Quot.sound']                                ACCEPT\n",
+      "t_simp       ['propext']                                   ACCEPT\n",
+      "axiom gate rejects exactly {native_decide, sorry, user axiom}; accepts constructive and classical-core proofs.\n"
+     ]
+    }
+   ],
+   "source": [
+    "import sys; sys.path.insert(0, os.path.abspath('.'))\n",
+    "from _lean_helper import Lean\n",
+    "L = Lean()\n",
+    "code = r'''\n",
+    "theorem k_comb (p q : Prop) : p → q → p := fun hp _ => hp\n",
+    "theorem peirce_law (p q : Prop) : ((p → q) → p) → p :=\n",
+    "  fun h => (Classical.em p).elim id (fun hn => h (fun hp => absurd hp hn))\n",
+    "theorem t_omega (a b : Nat) : a + b = b + a := by omega\n",
+    "theorem t_decide : 2^10 = 1024 := by decide\n",
+    "theorem t_native : 2^10 = 1024 := by native_decide\n",
+    "theorem t_sorry : (1:Nat) = 2 := sorry\n",
+    "axiom evil : False\n",
+    "theorem t_evil : (1:Nat) = 2 := evil.elim\n",
+    "theorem t_funext (f g : Nat → Nat) (h : ∀ x, f x = g x) : f = g := funext h\n",
+    "theorem t_simp (xs : List Nat) : (xs ++ []).length = xs.length := by simp\n",
+    "#print axioms k_comb\n",
+    "#print axioms peirce_law\n",
+    "#print axioms t_omega\n",
+    "#print axioms t_decide\n",
+    "#print axioms t_native\n",
+    "#print axioms t_sorry\n",
+    "#print axioms t_evil\n",
+    "#print axioms t_funext\n",
+    "#print axioms t_simp\n",
+    "'''\n",
+    "r = L.run(code, keep=True)\n",
+    "import re\n",
+    "axioms = {}\n",
+    "for m in r['messages']:\n",
+    "    mm = re.match(r\"'(\\w+)' (?:does not depend on any axioms|depends on axioms: \\[(.*)\\])\", m['data'])\n",
+    "    if mm: axioms[mm.group(1)] = sorted(x.strip() for x in (mm.group(2) or '').split(',') if x.strip())\n",
+    "ALLOWED = {'propext', 'Quot.sound', 'Classical.choice'}\n",
+    "print(f'{\"theorem\":12s} {\"axioms\":45s} gate')\n",
+    "for k_, v in axioms.items():\n",
+    "    print(f'{k_:12s} {str(v):45s} {\"ACCEPT\" if set(v) <= ALLOWED else \"REJECT\"}')\n",
+    "expected_reject = {'t_native', 't_sorry', 't_evil'}\n",
+    "assert {k_ for k_, v in axioms.items() if not set(v) <= ALLOWED} == expected_reject\n",
+    "assert axioms['k_comb'] == []\n",
+    "assert set(axioms['peirce_law']) == {'Classical.choice', 'propext', 'Quot.sound'}\n",
+    "print('axiom gate rejects exactly {native_decide, sorry, user axiom}; accepts constructive and classical-core proofs.')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3634e461",
+   "metadata": {},
+   "source": [
+    "### Tactics are untrusted; the term is what is checked\n",
+    "`by` blocks are programs that *build* a term. We can print it. A wrong tactic cannot make the kernel accept a wrong term."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "db937b1f",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:02.123377Z",
+     "iopub.status.busy": "2026-10-14T09:54:02.123216Z",
+     "iopub.status.idle": "2026-10-14T09:54:02.185056Z",
+     "shell.execute_reply": "2026-10-14T09:54:02.184044Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "`by omega` elaborates to a proof TERM of 64 pretty-printed lines (a certificate checked by the kernel); a short one:\n",
+      "[info] line 3: theorem t_by : ∀ (p q : Prop), p ∧ q → q ∧ p :=\n",
+      "fun p q h =>\n",
+      "  match h with\n",
+      "  | ⟨a, b⟩ => ⟨b, a⟩\n",
+      "--- ill-typed term: the kernel/elaborator refuses ---\n",
+      "[error] line 1: type mismatch\n",
+      "  Eq.refl\n",
+      "has type\n",
+      "  ∀ (a : ?m.2096), a = a : Prop\n",
+      "but is expected to have type\n",
+      "  1 = 2 : Prop\n",
+      "--- proof irrelevance: any two proofs of a Prop are definitionally equal ---\n",
+      "[ok] no messages\n",
+      "--- sorryAx inhabits every type, including False ---\n",
+      "[info] line 1: axiom sorryAx.{u} : (α : Sort u) → optParam Bool false → α\n",
+      "[warning] line 2: declaration uses 'sorry'\n",
+      "[info] line 3: 'anything' depends on axioms: [sorryAx]\n",
+      "[sorries] ⊢ False\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "{'sorries': [{'proofState': 1,\n",
+       "   'pos': {'line': 2, 'column': 28},\n",
+       "   'goal': '⊢ False',\n",
+       "   'endPos': {'line': 2, 'column': 33}}],\n",
+       " 'messages': [{'severity': 'info',\n",
+       "   'pos': {'line': 1, 'column': 0},\n",
+       "   'endPos': {'line': 1, 'column': 6},\n",
+       "   'data': 'axiom sorryAx.{u} : (α : Sort u) → optParam Bool false → α'},\n",
+       "  {'severity': 'warning',\n",
+       "   'pos': {'line': 2, 'column': 8},\n",
+       "   'endPos': {'line': 2, 'column': 16},\n",
+       "   'data': \"declaration uses 'sorry'\"},\n",
+       "  {'severity': 'info',\n",
+       "   'pos': {'line': 3, 'column': 0},\n",
+       "   'endPos': {'line': 3, 'column': 6},\n",
+       "   'data': \"'anything' depends on axioms: [sorryAx]\"}],\n",
+       " 'env': 7}"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "rr = L.run(\"#print t_omega\")\n",
+    "n_lines = sum(m['data'].count('\\n') + 1 for m in rr['messages'])\n",
+    "print(f'`by omega` elaborates to a proof TERM of {n_lines} pretty-printed lines (a certificate checked by the kernel); a short one:')\n",
+    "L.show(r'''\n",
+    "theorem t_by (p q : Prop) : p ∧ q → q ∧ p := by intro ⟨a, b⟩; exact ⟨b, a⟩\n",
+    "#print t_by\n",
+    "''')\n",
+    "print('--- ill-typed term: the kernel/elaborator refuses ---')\n",
+    "L.show(\"example : (1:Nat) = 2 := ⟨⟩\")\n",
+    "print('--- proof irrelevance: any two proofs of a Prop are definitionally equal ---')\n",
+    "L.show(\"example (p : Prop) (h1 h2 : p) : h1 = h2 := rfl\")\n",
+    "print('--- sorryAx inhabits every type, including False ---')\n",
+    "L.show(\"#print sorryAx\\ntheorem anything : False := sorry\\n#print axioms anything\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c0b25fdd",
+   "metadata": {},
+   "source": [
+    "### Why `sorryAx` / `native_decide` / extra axioms destroy a reward\n",
+    "Let `R(c) = 1[kernel accepts c and axioms(c) is within A]`. With A = {propext, Quot.sound, Classical.choice}, `R(c)=1`\n",
+    "implies the statement is true in every model of Lean's type theory plus those axioms (soundness of the kernel; consistency of A).\n",
+    "If `sorryAx` (or any axiom of type `False`) is allowed, `R(c) = 1` for **every** statement, so `P(R=1 | statement false) = 1`\n",
+    "and the reward carries zero information: an RL policy will find it. Below: acceptance probability of the trivial exploit under three gates."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "52d6fcf1",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:02.186945Z",
+     "iopub.status.busy": "2026-10-14T09:54:02.186790Z",
+     "iopub.status.idle": "2026-10-14T09:54:02.472246Z",
+     "shell.execute_reply": "2026-10-14T09:54:02.471199Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "false statement        axioms of the `sorry` proof strict-gate  no-axiom-gate(only compile)\n",
+      "(1:Nat) = 2            ['sorryAx']            REJECT       ACCEPT\n",
+      "∀ n : Nat, n < 3       ['sorryAx']            REJECT       ACCEPT\n",
+      "False                  ['sorryAx']            REJECT       ACCEPT\n"
+     ]
+    }
+   ],
+   "source": [
+    "statements_false = [\"(1:Nat) = 2\", \"∀ n : Nat, n < 3\", \"False\"]\n",
+    "rows = []\n",
+    "for st in statements_false:\n",
+    "    rr = L.run(f\"theorem probe : {st} := sorry\\n#print axioms probe\")\n",
+    "    ax = []\n",
+    "    for m in rr['messages']:\n",
+    "        mm = re.match(r\"'probe' depends on axioms: \\[(.*)\\]\", m['data'])\n",
+    "        if mm: ax = [x.strip() for x in mm.group(1).split(',')]\n",
+    "    rows.append((st, ax, 'ACCEPT' if set(ax) <= ALLOWED else 'REJECT', 'ACCEPT'))\n",
+    "print(f'{\"false statement\":22s} {\"axioms of the `sorry` proof\":22s} strict-gate  no-axiom-gate(only compile)')\n",
+    "for st, ax, g1, g2 in rows: print(f'{st:22s} {str(ax):22s} {g1:11s}  {g2}')\n",
+    "assert all(r_[2] == 'REJECT' for r_ in rows)\n",
+    "L.close()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  },
+  "m3v": {
+   "patched": "0301"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

code/notebooks/0302-dependent-types.ipynb

@@ -0,0 +1,762 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "00a22d32",
+   "metadata": {},
+   "source": [
+    "# 0302 - Dependent type theory: the minimum for reading Mathlib statements\n",
+    "\n",
+    "Scratch-pad for note `0302 - Dependent Types.pdf`. Naive first, better second, assert agreement.\n",
+    "\n",
+    "1. **Universe of a Pi type.** Naive: Python function `imax`-rule. Check against what Lean's kernel says for 10 types.\n",
+    "2. **Decidability.** Naive: Python brute force over the bound. Efficient: Lean `decide` (kernel/VM evaluation of the\n",
+    "   `Decidable` instance). Assert the Booleans agree on 12 bounded statements; show where no instance exists.\n",
+    "3. **Ten miniF2F statements decoded.** Lean prints the elaborated Pi-telescope; Python (exact `Fraction`/big-int arithmetic or\n",
+    "   fp64) shows the hypotheses are satisfiable (non-vacuous) and the conclusion holds at witnesses.\n",
+    "4. **Elaboration traps** (truncated subtraction, floor division, casts, `autoImplicit`), Lean vs a Python model of ℕ.\n",
+    "\n",
+    "No GPU (`CUDA_VISIBLE_DEVICES` cleared). Lean: ONE worker through `align.verify.lean`, nice 15."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "fa7036cc",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:20.065314Z",
+     "iopub.status.busy": "2026-10-14T09:54:20.065161Z",
+     "iopub.status.idle": "2026-10-14T09:54:24.370858Z",
+     "shell.execute_reply": "2026-10-14T09:54:24.370025Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
+    "import sys, re, json, math, random, itertools, time\n",
+    "from fractions import Fraction\n",
+    "sys.path.insert(0, os.path.abspath('.'))\n",
+    "from _lean_helper import Lean, HEADER\n",
+    "MINIF2F_HEADER = \"import Mathlib\\nimport Aesop\\n\\nset_option maxHeartbeats 400000\\n\\nopen BigOperators Real Nat Topology Rat\\n\"\n",
+    "L = Lean(MINIF2F_HEADER)\n",
+    "\n",
+    "def lean_msgs(code, keep=False, timeout=300):\n",
+    "    r = L.run(code, timeout=timeout, keep=keep)\n",
+    "    return r.get('messages', []), r"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6bd9c629",
+   "metadata": {},
+   "source": [
+    "## 1. The sort of a Pi type: `Sort (imax u v)`\n",
+    "Rule: if `A : Sort u` and `B x : Sort v` for `x : A` then `(x : A) → B x : Sort (imax u v)` where\n",
+    "`imax u 0 = 0` and `imax u v = max u v` for `v > 0`. So a Pi into `Prop` is a `Prop` whatever the domain (impredicativity)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "a0d26ffc",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:24.373206Z",
+     "iopub.status.busy": "2026-10-14T09:54:24.373041Z",
+     "iopub.status.idle": "2026-10-14T09:54:24.406046Z",
+     "shell.execute_reply": "2026-10-14T09:54:24.405001Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Pi type                      imax rule  Lean\n",
+      "∀ n : ℕ, n = n               Prop       Prop\n",
+      "ℕ → ℕ                        Type       Type\n",
+      "(n : ℕ) → Fin n              Type       Type\n",
+      "∀ p : Prop, p → p            Prop       Prop\n",
+      "(α : Type) → α → α           Type 1     Type 1\n",
+      "(α : Type) → Prop            Type 1     Type 1\n",
+      "Prop → Prop                  Type       Type\n",
+      "∀ α : Type, α = α            Prop       Prop\n",
+      "(α : Type 1) → α → α         Type 2     Type 2\n",
+      "ℕ → Prop                     Type       Type\n",
+      "imax rule agrees with the kernel on all 10 types.\n"
+     ]
+    }
+   ],
+   "source": [
+    "def imax(u, v): return 0 if v == 0 else max(u, v)\n",
+    "def sort_name(l): return 'Prop' if l == 0 else ('Type' if l == 1 else f'Type {l-1}')\n",
+    "\n",
+    "# (lean source of the Pi type, level u of the domain's sort, level v of the codomain's sort)\n",
+    "PI_TABLE = [\n",
+    "  (\"∀ n : ℕ, n = n\",              1, 0),   # ℕ : Type = Sort 1 ; n = n : Prop = Sort 0\n",
+    "  (\"ℕ → ℕ\",                       1, 1),\n",
+    "  (\"(n : ℕ) → Fin n\",             1, 1),\n",
+    "  (\"∀ p : Prop, p → p\",           1, 0),   # Prop : Type = Sort 1 ; p → p : Prop\n",
+    "  (\"(α : Type) → α → α\",          2, 1),   # Type : Type 1 = Sort 2 ; α → α : Type = Sort 1\n",
+    "  (\"(α : Type) → Prop\",           2, 1),   # Prop : Type = Sort 1\n",
+    "  (\"Prop → Prop\",                 1, 1),\n",
+    "  (\"∀ α : Type, α = α\",           2, 0),\n",
+    "  (\"(α : Type 1) → α → α\",        3, 2),   # Type 1 : Sort 3 ; α → α : Sort 2\n",
+    "  (\"ℕ → Prop\",                    1, 1),\n",
+    "]\n",
+    "rows = []\n",
+    "for src, u, v in PI_TABLE:\n",
+    "    msgs, _ = lean_msgs(f\"#check ({src})\")\n",
+    "    data = msgs[0]['data']\n",
+    "    m = re.search(r\": (Prop|Type\\*?|Type \\d+|Sort [^\\n]*)\\s*$\", data.strip().split('\\n')[-1])\n",
+    "    lean_sort = m.group(1) if m else data\n",
+    "    pred = sort_name(imax(u, v))\n",
+    "    rows.append((src, pred, lean_sort))\n",
+    "    assert pred == lean_sort, (src, pred, lean_sort)\n",
+    "print(f'{\"Pi type\":28s} {\"imax rule\":10s} Lean')\n",
+    "for r_ in rows: print(f'{r_[0]:28s} {r_[1]:10s} {r_[2]}')\n",
+    "print('imax rule agrees with the kernel on all', len(rows), 'types.')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f2877b3a",
+   "metadata": {},
+   "source": [
+    "## 2. Decidability: `Decidable p` is data; `decide` runs it\n",
+    "`class Decidable (p : Prop) : Type | isFalse (h : ¬p) | isTrue (h : p)`. `decide p` is `match inst with | isTrue _ => true | isFalse _ => false`.\n",
+    "Soundness: `of_decide_eq_true : decide p = true → p` (cases on the instance). The tactic builds `of_decide_eq_true (Eq.refl true)`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "f6a87131",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:24.407763Z",
+     "iopub.status.busy": "2026-10-14T09:54:24.407583Z",
+     "iopub.status.idle": "2026-10-14T09:54:24.427808Z",
+     "shell.execute_reply": "2026-10-14T09:54:24.427297Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "@of_decide_eq_true : ∀ {p : Prop} [inst : Decidable p], decide p = true → p\n",
+      "inductive Decidable : Prop → Type\n",
+      "number of parameters: 1\n",
+      "constructors:\n",
+      "Decidable.isFalse : {p : Prop} → ¬p → Decidable p\n",
+      "Decidable.isTrue : {p : Prop} → p → Decidable p\n",
+      "theorem two_plus_two : 2 + 2 = 4 :=\n",
+      "of_decide_eq_true (Eq.refl true)\n",
+      "'two_plus_two' does not depend on any axioms\n"
+     ]
+    }
+   ],
+   "source": [
+    "msgs, _ = lean_msgs(\"#check @of_decide_eq_true\\n#print Decidable\\ntheorem two_plus_two : 2 + 2 = 4 := by decide\\n#print two_plus_two\\n#print axioms two_plus_two\", keep=True)\n",
+    "for m in msgs: print(m['data'])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "357adce9",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:24.433916Z",
+     "iopub.status.busy": "2026-10-14T09:54:24.433656Z",
+     "iopub.status.idle": "2026-10-14T09:54:25.615319Z",
+     "shell.execute_reply": "2026-10-14T09:54:25.614598Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "--- Nat.Prime is decidable only via a *chosen* algorithm; two instances exist (the default is chosen by priority) ---\n",
+      "decidablePrime 7\n",
+      "decidablePrime1 7\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "python  lean   statement\n",
+      "True    true   ∀ n < 50, n * n % 4 ≠ 2\n",
+      "True    true   ∃ n < 100, n * n = 2209\n",
+      "True    true   ∀ a < 20, ∀ b < 20, a * b = b * a\n",
+      "True    true   ∀ n < (40 : ℕ), Nat.Prime (n * n + n + 41)\n",
+      "False   false  ∀ n < (41 : ℕ), Nat.Prime (n * n + n + 41)\n",
+      "True    true   ∀ n < 100, n ^ 3 % 7 = 0 ∨ n ^ 3 % 7 = 1 ∨ n ^ 3 % 7 = 6\n",
+      "False   false  ∃ x : Fin 61, ∃ y : Fin 61, 4 * (((x : ℕ) : ℤ) - 30) ^ 3 - 7 * (((y : ℕ) : ℤ) - 30) ^ 3 = 2003\n",
+      "False   false  ∀ a < 10, ∀ b < 10, a - b + b = a\n",
+      "True    true   ∀ a < 10, ∀ b < 10, (a + b) / 2 * 2 ≤ a + b\n",
+      "False   false  ∃ n < 10, 18 ∣ 3740 + n ∧ n ≠ 4\n",
+      "True    true   ∀ n < 30, 2 ^ n % 3 = if n % 2 = 0 then 1 else 2\n",
+      "True    true   ∃ n < 100, n % 3 = 2 ∧ n % 5 = 3 ∧ n % 7 = 2\n",
+      "python brute force == Lean decide on 12/12\n"
+     ]
+    }
+   ],
+   "source": [
+    "# 12 bounded statements: (python oracle, lean statement)\n",
+    "def prime(n):\n",
+    "    return n >= 2 and all(n % d for d in range(2, int(n ** .5) + 1))\n",
+    "def nsub(a, b): return max(0, a - b)          # Lean ℕ subtraction is truncated\n",
+    "CASES = [\n",
+    " (lambda: all(n*n % 4 != 2 for n in range(50)),                         \"∀ n < 50, n * n % 4 ≠ 2\"),\n",
+    " (lambda: any(n*n == 2209 for n in range(100)),                         \"∃ n < 100, n * n = 2209\"),\n",
+    " (lambda: all(a*b == b*a for a in range(20) for b in range(20)),        \"∀ a < 20, ∀ b < 20, a * b = b * a\"),\n",
+    " (lambda: all(prime(n*n+n+41) for n in range(40)),                      \"∀ n < (40 : ℕ), Nat.Prime (n * n + n + 41)\"),\n",
+    " (lambda: all(prime(n*n+n+41) for n in range(41)),                      \"∀ n < (41 : ℕ), Nat.Prime (n * n + n + 41)\"),\n",
+    " (lambda: all(n**3 % 7 in (0, 1, 6) for n in range(100)),               \"∀ n < 100, n ^ 3 % 7 = 0 ∨ n ^ 3 % 7 = 1 ∨ n ^ 3 % 7 = 6\"),\n",
+    " (lambda: any(4*x**3 - 7*y**3 == 2003 for x in range(-30, 31) for y in range(-30, 31)),\n",
+    "                                                                        \"∃ x : Fin 61, ∃ y : Fin 61, 4 * (((x : ℕ) : ℤ) - 30) ^ 3 - 7 * (((y : ℕ) : ℤ) - 30) ^ 3 = 2003\"),\n",
+    " (lambda: all(nsub(a, b) + b == a for a in range(10) for b in range(10)),   \"∀ a < 10, ∀ b < 10, a - b + b = a\"),\n",
+    " (lambda: all((a + b) // 2 * 2 <= a + b for a in range(10) for b in range(10)), \"∀ a < 10, ∀ b < 10, (a + b) / 2 * 2 ≤ a + b\"),\n",
+    " (lambda: any((3740 + n) % 18 == 0 and n != 4 for n in range(10)),      \"∃ n < 10, 18 ∣ 3740 + n ∧ n ≠ 4\"),\n",
+    " (lambda: all(pow(2, n) % 3 == (1 if n % 2 == 0 else 2) for n in range(30)), \"∀ n < 30, 2 ^ n % 3 = if n % 2 = 0 then 1 else 2\"),\n",
+    " (lambda: any(n % 3 == 2 and n % 5 == 3 and n % 7 == 2 for n in range(100)), \"∃ n < 100, n % 3 = 2 ∧ n % 5 = 3 ∧ n % 7 = 2\"),\n",
+    "]\n",
+    "print('--- Nat.Prime is decidable only via a *chosen* algorithm; two instances exist (the default is chosen by priority) ---')\n",
+    "for m in lean_msgs(\"#synth Decidable (Nat.Prime 7)\")[0]: print(m['data'][:160])\n",
+    "lean_msgs(\"attribute [instance] Nat.decidablePrime1\", keep=True)      # opt in to the trial-division instance (needed below for bounded foralls)\n",
+    "for m in lean_msgs(\"#synth Decidable (Nat.Prime 7)\")[0]: print(m['data'][:160])\n",
+    "code = \"\\n\".join(f\"#eval decide ({s})\" for _, s in CASES)\n",
+    "msgs, _ = lean_msgs(code)\n",
+    "by_line = {m['pos']['line']: m['data'].strip() for m in msgs}\n",
+    "print(f'{\"python\":7s} {\"lean\":6s} statement')\n",
+    "agree = 0\n",
+    "for i, (fn, s) in enumerate(CASES, start=1):\n",
+    "    py = fn(); ln = by_line.get(i)\n",
+    "    assert ln in ('true', 'false'), (s, ln)\n",
+    "    assert (ln == 'true') == py, (s, py, ln)\n",
+    "    agree += 1\n",
+    "    print(f'{str(py):7s} {ln:6s} {s}')\n",
+    "print(f'python brute force == Lean decide on {agree}/{len(CASES)}')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "c58167d4",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:25.617088Z",
+     "iopub.status.busy": "2026-10-14T09:54:25.616927Z",
+     "iopub.status.idle": "2026-10-14T09:54:25.737634Z",
+     "shell.execute_reply": "2026-10-14T09:54:25.736825Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "--- elaboration order matters: the same statement without the type ascription has no instance (n : ?α is still a metavariable) ---\n",
+      "failed to synthesize\n",
+      "  Decidable (∀ n < 40, (n * n + n + 41).Prime)\n",
+      "use `set_option diagnostics true` to get diagnostic information\n",
+      "--- no instance: unbounded quantifier over ℕ ---\n",
+      "failed to synthesize\n",
+      "  Decidable (∀ (n : ℕ), n < 5 ∨ 5 ≤ n)\n",
+      "use `set_option diagnostics true` to get diagnostic information\n",
+      "--- no computable instance: equality on ℝ ---\n",
+      "tactic 'decide' failed for proposition\n",
+      "  2 + 2 = 4\n",
+      "since its 'Decidable' instance reduced to\n",
+      "  Classical.choice ⋯\n",
+      "rather than to the 'isTrue' constructor.\n",
+      "--- but a *closed* bounded one is fine, and the real statements with binders are not decidable props at all ---\n",
+      "decidableBallLT 10 fun n a => n * n < 100\n",
+      "failed to synthesize\n",
+      "  Decidable (∀ (x : ℝ), x ^ 2 ≥ 0)\n",
+      "use `set_option diagnostics true` to get diagnostic information\n"
+     ]
+    }
+   ],
+   "source": [
+    "print('--- elaboration order matters: the same statement without the type ascription has no instance (n : ?α is still a metavariable) ---')\n",
+    "for m in lean_msgs(\"#eval decide (∀ n < 40, Nat.Prime (n * n + n + 41))\")[0]: print(m['data'][:160])\n",
+    "print('--- no instance: unbounded quantifier over ℕ ---')\n",
+    "for m in lean_msgs(\"example : ∀ n : ℕ, n < 5 ∨ 5 ≤ n := by decide\")[0]: print(m['data'][:400])\n",
+    "print('--- no computable instance: equality on ℝ ---')\n",
+    "for m in lean_msgs(\"example : (2 : ℝ) + 2 = 4 := by decide\")[0]: print(m['data'][:400])\n",
+    "print('--- but a *closed* bounded one is fine, and the real statements with binders are not decidable props at all ---')\n",
+    "for m in lean_msgs(\"#synth Decidable (∀ n < 10, n * n < 100)\")[0]: print(m['data'][:200])\n",
+    "for m in lean_msgs(\"#synth Decidable (∀ x : ℝ, x ^ 2 ≥ 0)\")[0]: print(m['data'][:200])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c172fa47",
+   "metadata": {},
+   "source": [
+    "## 3. Ten miniF2F statements decoded\n",
+    "Source: `p2_runs/r1_baseline/minif2f.jsonl` (the 244 miniF2F test statements this repo evaluates on, Lean 4.9.0-rc1 + Mathlib pin).\n",
+    "For each: (a) Lean's own elaborated Pi-telescope (`#check @name`), its sort, and `decide`-ability of the closed ones;\n",
+    "(b) a Python witness check that the hypotheses are satisfiable and the conclusion holds there (non-vacuity), exact arithmetic where possible."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "f06bbc76",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:25.739310Z",
+     "iopub.status.busy": "2026-10-14T09:54:25.739151Z",
+     "iopub.status.idle": "2026-10-14T09:54:26.102413Z",
+     "shell.execute_reply": "2026-10-14T09:54:26.101621Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "====================================================================================================\n",
+      "mathd_algebra_478\n",
+      "mathd_algebra_478 : ∀ (b h v : ℝ), 0 < b ∧ 0 < h ∧ 0 < v → v = 1 / 3 * (b * h) → b = 30 → h = 13 / 2 → v = 65\n",
+      "binders=7 sort=Prop\n",
+      "====================================================================================================\n",
+      "numbertheory_4x3m7y3neq2003\n",
+      "numbertheory_4x3m7y3neq2003 : ∀ (x y : ℤ), 4 * x ^ 3 - 7 * y ^ 3 ≠ 2003\n",
+      "binders=2 sort=Prop\n",
+      "====================================================================================================\n",
+      "aime_1983_p1\n",
+      "aime_1983_p1 : ∀ (x y z w : ℕ), 1 < x ∧ 1 < y ∧ 1 < z → 0 ≤ w → (↑w).log / (↑x).log = 24 → (↑w).log / (↑y).log = 40 → (↑w).log / (↑x * ↑y * ↑z).log = 12 → (↑w).log / (↑z).log = 60\n",
+      "binders=9 sort=Prop\n",
+      "====================================================================================================\n",
+      "amc12_2001_p5\n",
+      "amc12_2001_p5 : (Finset.filter (fun x => ¬Even x) (Finset.range 10000)).prod id = 10000! / (2 ^ 5000 * 5000!)\n",
+      "binders=0 sort=Prop\n",
+      "====================================================================================================\n",
+      "mathd_numbertheory_3\n",
+      "mathd_numbertheory_3 : (∑ x ∈ Finset.range 10, (x + 1) ^ 2) % 10 = 5\n",
+      "binders=0 sort=Prop\n",
+      "====================================================================================================\n",
+      "imo_1969_p2\n",
+      "imo_1969_p2 : ∀ (m n : ℝ) (k : ℕ) (a : ℕ → ℝ) (y : ℝ → ℝ), 0 < k → (∀ (x : ℝ), y x = ∑ i ∈ Finset.range k, (a i + x).cos / 2 ^ i) → y m = 0 → y n = 0 → ∃ t, m - n = ↑t * π\n",
+      "binders=9 sort=Prop\n",
+      "====================================================================================================\n",
+      "mathd_algebra_209\n",
+      "mathd_algebra_209 : ∀ (σ : ℝ ≃ ℝ), σ.invFun 2 = 10 → σ.invFun 10 = 1 → σ.invFun 1 = 2 → σ.toFun (σ.toFun 10) = 1\n",
+      "binders=4 sort=Prop\n",
+      "====================================================================================================\n",
+      "mathd_numbertheory_1124\n",
+      "mathd_numbertheory_1124 : ∀ n ≤ 9, 18 ∣ 374 * 10 + n → n = 4\n",
+      "binders=3 sort=Prop\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "====================================================================================================\n",
+      "imo_1983_p6\n",
+      "imo_1983_p6 : ∀ (a b c : ℝ), 0 < a ∧ 0 < b ∧ 0 < c → c < a + b → b < a + c → a < b + c → 0 ≤ a ^ 2 * b * (a - b) + b ^ 2 * c * (b - c) + c ^ 2 * a * (c - a)\n",
+      "binders=7 sort=Prop\n",
+      "====================================================================================================\n",
+      "mathd_algebra_33\n",
+      "mathd_algebra_33 : ∀ (x y z : ℝ), x ≠ 0 → 2 * x = 5 * y → 7 * y = 10 * z → z / x = 7 / 25\n",
+      "binders=6 sort=Prop\n"
+     ]
+    }
+   ],
+   "source": [
+    "PICK = ['mathd_algebra_478', 'numbertheory_4x3m7y3neq2003', 'aime_1983_p1', 'amc12_2001_p5', 'mathd_numbertheory_3',\n",
+    "        'imo_1969_p2', 'mathd_algebra_209', 'mathd_numbertheory_1124', 'imo_1983_p6', 'mathd_algebra_33']\n",
+    "rows = {}\n",
+    "for l_ in open('/apps/development/rlhf/p2_runs/r1_baseline/minif2f.jsonl'):\n",
+    "    r_ = json.loads(l_)\n",
+    "    if r_['problem_id'] in PICK: rows[r_['problem_id']] = r_\n",
+    "assert set(rows) == set(PICK)\n",
+    "\n",
+    "SORT_ELAB = r'''\n",
+    "open Lean Meta Elab Command in\n",
+    "elab \"#sort_of \" id:ident : command => do\n",
+    "  let env ← getEnv\n",
+    "  match env.find? id.getId with\n",
+    "  | none => throwError \"unknown\"\n",
+    "  | some ci => liftTermElabM do\n",
+    "      let s ← inferType ci.type\n",
+    "      let nb := (← forallTelescope ci.type fun xs _ => pure xs.size)\n",
+    "      logInfo m!\"binders={nb} sort={← whnf s}\"\n",
+    "'''\n",
+    "L.run(SORT_ELAB, keep=True)\n",
+    "\n",
+    "decoded = {}\n",
+    "for pid in PICK:\n",
+    "    st = rows[pid]['formal_statement']\n",
+    "    thm = st[st.index('/--'):]                                   # docstring + theorem ... := by sorry\n",
+    "    rest = thm\n",
+    "    msgs, r = lean_msgs(rest + f\"\\n#check @{pid}\\n#sort_of {pid}\", timeout=300)\n",
+    "    out = [m['data'] for m in msgs if m['severity'] == 'info']\n",
+    "    errs = [m['data'] for m in msgs if m['severity'] == 'error']\n",
+    "    assert not errs, errs\n",
+    "    decoded[pid] = out\n",
+    "    print('=' * 100); print(pid)\n",
+    "    print(re.sub(r'\\n\\s+', ' ', out[0]))\n",
+    "    print(out[1])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2d1b76ed",
+   "metadata": {},
+   "source": [
+    "Closed statements (no binders) are plain `Prop`s with a `Decidable` instance: Lean can *evaluate* them."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "852923ed",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:26.104415Z",
+     "iopub.status.busy": "2026-10-14T09:54:26.104258Z",
+     "iopub.status.idle": "2026-10-14T09:54:26.288412Z",
+     "shell.execute_reply": "2026-10-14T09:54:26.287495Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "amc12_2001_p5 -> ['instDecidableEqNat ((Finset.filter (fun x => ¬Even x) (Finset.range 10000)).prod', 'true']\n",
+      "mathd_numbertheory_3 -> ['instDecidableEqNat ((∑ x ∈ Finset.range 10, (x + 1) ^ 2) % 10) 5', 'true']\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[\"'nt3' depends on axioms: [propext, Quot.sound]\"]\n"
+     ]
+    }
+   ],
+   "source": [
+    "for pid in ['amc12_2001_p5', 'mathd_numbertheory_3']:\n",
+    "    st = rows[pid]['formal_statement']\n",
+    "    body = st.split(f'theorem {pid} :')[1].split(':= by sorry')[0].strip()\n",
+    "    msgs, _ = lean_msgs(f\"#synth Decidable ({body})\\n#eval decide ({body})\", timeout=300)\n",
+    "    print(pid, '->', [m['data'].strip()[:80] for m in msgs])\n",
+    "msgs, _ = lean_msgs(\"theorem nt3 : (∑ x in Finset.range 10, (x + 1) ^ 2) % 10 = 5 := by decide\\n#print axioms nt3\")\n",
+    "print([m['data'] for m in msgs])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "b59c7854",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:26.290114Z",
+     "iopub.status.busy": "2026-10-14T09:54:26.289954Z",
+     "iopub.status.idle": "2026-10-14T09:54:26.591691Z",
+     "shell.execute_reply": "2026-10-14T09:54:26.590442Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "mathd_algebra_478              OK  witness b=30,h=13/2 => v=65 exactly\n",
+      "numbertheory_4x3m7y3neq2003    OK  cubes mod 7 = [0, 1, 6]; 4*cubes mod 7 = [0, 3, 4]; 2003 mod 7 = 1 not attained; box [-80,80]^2 empty\n",
+      "aime_1983_p1                   OK  fp64 max hypothesis residual 0.00e+00, conclusion residual 0.00e+00; exact: L/(L/12-L/24-L/40)=60\n",
+      "amc12_2001_p5                  OK  product of odd numbers < 10000 == 10000!/(2^5000 * 5000!) as big integers; division exact\n",
+      "mathd_numbertheory_3           OK  sum = 385, 385 mod 10 = 5\n",
+      "imo_1969_p2                    OK  200 random (k,a): y(x)=A cos x - B sin x; zeros m,n: max |y(m)|,|y(n)|,|(m-n)/π - t| = 3.7e-15; |(A,B)| >= 2^(1-k) > 0 in all 200 runs\n",
+      "mathd_algebra_209              OK  σ = {10: 2, 1: 10, 2: 1}; σ(σ(10)) = 1\n",
+      "mathd_numbertheory_1124        OK  solutions n in 0..9 with 18 | 3740+n: [4]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "imo_1983_p6                    OK  min over 2e5 random triangles of E/a^4 = 6.588e-08 (>= 0); a=b=c gives exactly 0\n",
+      "mathd_algebra_33               OK  x=1,y=2/5,z=7/25: z/x = 7/25 exactly\n",
+      "all 10 statements: hypotheses satisfiable (non-vacuous) and conclusion verified at witnesses / exhaustively.\n"
+     ]
+    }
+   ],
+   "source": [
+    "# ---- (b) Python witness checks, one per statement ----\n",
+    "def check_478():\n",
+    "    b, h = Fraction(30), Fraction(13, 2)\n",
+    "    v = Fraction(1, 3) * (b * h)\n",
+    "    assert 0 < b and 0 < h and 0 < v          # hypotheses satisfiable (h₀)\n",
+    "    assert v == 65                           # conclusion\n",
+    "    return 'witness b=30,h=13/2 => v=65 exactly'\n",
+    "\n",
+    "def check_4x3m7y3():\n",
+    "    cubes = {pow(n, 3, 7) for n in range(7)}\n",
+    "    resid = {(4 * c) % 7 for c in cubes}\n",
+    "    assert 2003 % 7 == 1 and 1 not in resid  # the real proof: 4x^3 - 7y^3 ≡ 4x^3 (mod 7)\n",
+    "    assert not any(4*x**3 - 7*y**3 == 2003 for x in range(-80, 81) for y in range(-80, 81))\n",
+    "    return f'cubes mod 7 = {sorted(cubes)}; 4*cubes mod 7 = {sorted(resid)}; 2003 mod 7 = 1 not attained; box [-80,80]^2 empty'\n",
+    "\n",
+    "def check_aime():\n",
+    "    # w = 2^120, x = 2^5, y = 2^3, z = 2^2 : satisfies all three hypotheses; conclusion log_z w = 60\n",
+    "    x, y, z, w = 2**5, 2**3, 2**2, 2**120\n",
+    "    errs = [abs(math.log(w)/math.log(x) - 24), abs(math.log(w)/math.log(y) - 40), abs(math.log(w)/math.log(x*y*z) - 12)]\n",
+    "    concl = abs(math.log(w)/math.log(z) - 60)\n",
+    "    assert max(errs) < 1e-12 and concl < 1e-12\n",
+    "    Lw = Fraction(120); lz = Lw/12 - Lw/24 - Lw/40      # log z = log(xyz) - log x - log y\n",
+    "    assert Lw / lz == 60                                  # exact derivation\n",
+    "    return f'fp64 max hypothesis residual {max(errs):.2e}, conclusion residual {concl:.2e}; exact: L/(L/12-L/24-L/40)=60'\n",
+    "\n",
+    "def check_amc():\n",
+    "    prod = 1\n",
+    "    for x in range(10000):\n",
+    "        if x % 2 == 1: prod *= x\n",
+    "    rhs_num = math.factorial(10000); rhs_den = 2**5000 * math.factorial(5000)\n",
+    "    assert rhs_num % rhs_den == 0               # the ℕ-division is exact, so truncation plays no role\n",
+    "    assert prod == rhs_num // rhs_den\n",
+    "    return 'product of odd numbers < 10000 == 10000!/(2^5000 * 5000!) as big integers; division exact'\n",
+    "\n",
+    "def check_nt3():\n",
+    "    s = sum((x + 1) ** 2 for x in range(10)); assert s == 385 and s % 10 == 5\n",
+    "    return 'sum = 385, 385 mod 10 = 5'\n",
+    "\n",
+    "def check_imo69():\n",
+    "    random.seed(1); worst = 0.0\n",
+    "    for _ in range(200):\n",
+    "        k = random.randint(1, 8); a = [random.uniform(-10, 10) for _ in range(k)]\n",
+    "        A = sum(math.cos(a[i]) / 2**i for i in range(k)); B = sum(math.sin(a[i]) / 2**i for i in range(k))\n",
+    "        x0 = math.atan2(A, B)                      # y(x) = A cos x - B sin x = 0  <=>  x = atan2(A,B) + tπ\n",
+    "        y = lambda x: sum(math.cos(a[i] + x) / 2**i for i in range(k))\n",
+    "        m, n = x0 + 3 * math.pi, x0 - 2 * math.pi\n",
+    "        worst = max(worst, abs(y(m)), abs(y(n)), abs((m - n) / math.pi - 5))\n",
+    "        assert math.hypot(A, B) >= 2.0 ** (1 - k) - 1e-12     # |sum_i e^{i a_i} 2^-i| >= 1 - sum_{i=1}^{k-1} 2^-i = 2^(1-k) > 0 (FINITE sum)\n",
+    "    assert worst < 1e-12\n",
+    "    return f'200 random (k,a): y(x)=A cos x - B sin x; zeros m,n: max |y(m)|,|y(n)|,|(m-n)/π - t| = {worst:.1e}; |(A,B)| >= 2^(1-k) > 0 in all 200 runs'\n",
+    "\n",
+    "def check_209():\n",
+    "    inv = {2: 10, 10: 1, 1: 2}                    # σ.2 = invFun\n",
+    "    sigma = {v: k for k, v in inv.items()}\n",
+    "    assert sigma[sigma[10]] == 1\n",
+    "    return f'σ = {sigma}; σ(σ(10)) = {sigma[sigma[10]]}'\n",
+    "\n",
+    "def check_1124():\n",
+    "    sol = [n for n in range(10) if (374 * 10 + n) % 18 == 0]\n",
+    "    assert sol == [4]\n",
+    "    return f'solutions n in 0..9 with 18 | 3740+n: {sol}'\n",
+    "\n",
+    "def check_1983():\n",
+    "    random.seed(2); worst = float('inf')\n",
+    "    for _ in range(200000):\n",
+    "        x, y, z = (random.uniform(1e-3, 10) for _ in range(3))      # Ravi: a=y+z, b=z+x, c=x+y  => triangle inequalities hold\n",
+    "        a, b, c = y + z, z + x, x + y\n",
+    "        val = a*a*b*(a-b) + b*b*c*(b-c) + c*c*a*(c-a)\n",
+    "        worst = min(worst, val / (a**4))                            # scale-free\n",
+    "    assert worst > -1e-12\n",
+    "    t = Fraction(7); assert t*t*t*(t-t) * 3 == 0\n",
+    "    return f'min over 2e5 random triangles of E/a^4 = {worst:.3e} (>= 0); a=b=c gives exactly 0'\n",
+    "\n",
+    "def check_33():\n",
+    "    x = Fraction(1); y = 2 * x / 5; z = 7 * y / 10\n",
+    "    assert x != 0 and 2*x == 5*y and 7*y == 10*z and z / x == Fraction(7, 25)\n",
+    "    return f'x=1,y={y},z={z}: z/x = 7/25 exactly'\n",
+    "\n",
+    "CHECKS = {'mathd_algebra_478': check_478, 'numbertheory_4x3m7y3neq2003': check_4x3m7y3, 'aime_1983_p1': check_aime,\n",
+    "          'amc12_2001_p5': check_amc, 'mathd_numbertheory_3': check_nt3, 'imo_1969_p2': check_imo69,\n",
+    "          'mathd_algebra_209': check_209, 'mathd_numbertheory_1124': check_1124, 'imo_1983_p6': check_1983,\n",
+    "          'mathd_algebra_33': check_33}\n",
+    "for pid in PICK:\n",
+    "    print(f'{pid:30s} OK  {CHECKS[pid]()}')\n",
+    "print('all 10 statements: hypotheses satisfiable (non-vacuous) and conclusion verified at witnesses / exhaustively.')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4fadce18",
+   "metadata": {},
+   "source": [
+    "## 3b. Lemma checks (section 3.4 of the note)\n",
+    "`(a - b) + b = max a b` on ℕ (exhaustive in a box AND a Lean proof); the Euclidean division law `a = b*(a/b) + a%b, 0 <= a%b < |b|` for Lean's ℤ (exhaustive grid evaluated by Lean); and `x / 0 = 0` on ℝ."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "165b7f70",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:26.593548Z",
+     "iopub.status.busy": "2026-10-14T09:54:26.593364Z",
+     "iopub.status.idle": "2026-10-14T09:54:27.071951Z",
+     "shell.execute_reply": "2026-10-14T09:54:27.071277Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Lean: a - b + b = max a b proved; 'tsub_add_max' depends on axioms: [propext, Quot.sound] | and the naive \"a - b + b = a\" is refuted at (0,1)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Lean ℤ: b*(a/b) + a%b = a and 0 <= a%b < |b| on 1464 grid points; Lean (a/b, a%b) == the unique Euclidean (q,r) on all of them; x/0 = 0, x%0 = x\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Lean ℝ: ∀ x, x / 0 = 0 proved (div_zero); 'real_div_zero' depends on axioms: [propext, Classical.choice, Quot.sound]\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "# (i) truncated subtraction: (a - b) + b = max a b\n",
+    "assert all(max(0, a - b) + b == max(a, b) for a in range(61) for b in range(61))\n",
+    "msgs, _ = lean_msgs(\"\"\"\n",
+    "theorem tsub_add_max (a b : ℕ) : a - b + b = max a b := by omega\n",
+    "#print axioms tsub_add_max\n",
+    "theorem tsub_add_false : ¬ (∀ a b : ℕ, a - b + b = a) := fun h => by have := h 0 1; simp at this\n",
+    "\"\"\")\n",
+    "assert not [m for m in msgs if m['severity'] == 'error'], msgs\n",
+    "print('Lean: a - b + b = max a b proved;', [m['data'] for m in msgs if m['severity'] == 'info'][0], '| and the naive \"a - b + b = a\" is refuted at (0,1)')\n",
+    "# (ii) Euclidean division law for Int, exhaustive on a grid, evaluated by Lean\n",
+    "msgs, _ = lean_msgs(\"\"\"\n",
+    "#eval (List.range 61).all fun ai => (List.range 25).all fun bi =>\n",
+    "  let a : Int := (ai : Int) - 30\n",
+    "  let b : Int := (bi : Int) - 12\n",
+    "  b == 0 || ((b * (a / b) + a % b == a) && decide (0 ≤ a % b) && decide (a % b < (b.natAbs : Int)))\n",
+    "#eval (List.range 61).all fun ai => (List.range 25).all fun bi => ((ai : Int) - 30) / 0 == 0 && ((ai : Int) - 30) % 0 == (ai : Int) - 30\n",
+    "\"\"\")\n",
+    "res = [m['data'].strip() for m in msgs if m['severity'] == 'info']\n",
+    "assert res == ['true', 'true'], (res, msgs)\n",
+    "# the unique (q, r) with a = b q + r, 0 <= r < |b| (Python model) must equal Lean's / and %\n",
+    "def euclid(a, b):\n",
+    "    r = a % abs(b); q = (a - r) // b; return q, r\n",
+    "lean_vals = lean_msgs(\"#eval (List.range 61).bind fun ai => (List.range 25).filterMap fun bi => let a : Int := (ai : Int) - 30; let b : Int := (bi : Int) - 12; if b == 0 then none else some (a / b, a % b)\")[0]\n",
+    "import ast\n",
+    "lv = ast.literal_eval(lean_vals[0]['data'].replace('(', '(').strip()) if lean_vals and lean_vals[0]['severity'] == 'info' else None\n",
+    "pv = [euclid(ai - 30, bi - 12) for ai in range(61) for bi in range(25) if bi != 12]\n",
+    "assert lv is not None and [tuple(x) for x in lv] == pv, 'Lean Int / and % differ from the Euclidean (q, r) model'\n",
+    "print(f'Lean ℤ: b*(a/b) + a%b = a and 0 <= a%b < |b| on {len(pv)} grid points; Lean (a/b, a%b) == the unique Euclidean (q,r) on all of them; x/0 = 0, x%0 = x')\n",
+    "# (iii) x / 0 = 0 on the reals\n",
+    "msgs, _ = lean_msgs(\"theorem real_div_zero : ∀ x : ℝ, x / 0 = 0 := fun x => div_zero x\\n#print axioms real_div_zero\")\n",
+    "assert not [m for m in msgs if m['severity'] == 'error'], msgs\n",
+    "print('Lean ℝ: ∀ x, x / 0 = 0 proved (div_zero);', [m['data'] for m in msgs if m['severity'] == 'info'][0])\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8d2c076e",
+   "metadata": {},
+   "source": [
+    "## 4. Elaboration traps (where a *correct-looking* statement means something else)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "85822ae4",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:27.074232Z",
+     "iopub.status.busy": "2026-10-14T09:54:27.074069Z",
+     "iopub.status.idle": "2026-10-14T09:54:27.447944Z",
+     "shell.execute_reply": "2026-10-14T09:54:27.447329Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(5 : ℕ) - 7                  = 0\n",
+      "(7 : ℕ) / 2                  = 3\n",
+      "(7 : ℤ) / (-2)               = -3\n",
+      "(-7 : ℤ) / 2                 = -4\n",
+      "(-7 : ℤ) % 3                 = 2\n",
+      "(((5 : ℕ) - 7 : ℕ) : ℤ)      = 0\n",
+      "((5 : ℕ) : ℤ) - 7            = -2\n",
+      "((10 : ℕ) - 3) / 2           = 3\n",
+      "(2 : ℕ) ^ 3 ^ 2              = 512\n",
+      "Lean 7/(-2) = -3  and (-7)/2 = -4 (Euclidean: remainder is always >= 0)\n",
+      "--- autoImplicit: a typo silently becomes a new universally-quantified variable ---\n",
+      "warning declaration uses 'sorry'\n",
+      "info @t_typo : ∀ {xx : ℕ} (x : ℕ), x + 1 > xx\n",
+      "--- the same with autoImplicit off: an error, as it should be ---\n",
+      "error unknown identifier 'xx'\n"
+     ]
+    }
+   ],
+   "source": [
+    "TRAPS = [\n",
+    " (\"(5 : ℕ) - 7\",           lambda: nsub(5, 7),            \"0\"),\n",
+    " (\"(7 : ℕ) / 2\",           lambda: 7 // 2,                \"3\"),\n",
+    " (\"(7 : ℤ) / (-2)\",        lambda: 'ediv',                None),\n",
+    " (\"(-7 : ℤ) / 2\",          lambda: 'ediv',                None),\n",
+    " (\"(-7 : ℤ) % 3\",          lambda: (-7) % 3,              \"2\"),\n",
+    " (\"(((5 : ℕ) - 7 : ℕ) : ℤ)\", lambda: nsub(5, 7),          \"0\"),\n",
+    " (\"((5 : ℕ) : ℤ) - 7\",     lambda: 5 - 7,                 \"-2\"),\n",
+    " (\"((10 : ℕ) - 3) / 2\",    lambda: (10 - 3) // 2,         \"3\"),\n",
+    " (\"(2 : ℕ) ^ 3 ^ 2\",       lambda: 2 ** (3 ** 2),         \"512\"),\n",
+    "]\n",
+    "msgs, _ = lean_msgs(\"\\n\".join(f\"#eval {s}\" for s, _, _ in TRAPS))\n",
+    "by_line = {m['pos']['line']: m['data'].strip() for m in msgs}\n",
+    "for i, (s, fn, expect) in enumerate(TRAPS, start=1):\n",
+    "    got = by_line[i]\n",
+    "    if expect is not None:\n",
+    "        assert got == expect, (s, got, expect)\n",
+    "        py = fn(); assert str(py) == got, (s, py, got)\n",
+    "    print(f'{s:28s} = {got}')\n",
+    "# Lean's ℤ `/` is Euclidean-style `Int.div` (T-rounding is `Int.tdiv`): -7/2 = -4, 7/(-2) = -3 ; Python's // floors: -7//2 = -4, 7//(-2) = -4\n",
+    "print('Lean 7/(-2) =', by_line[3], ' and (-7)/2 =', by_line[4], '(Euclidean: remainder is always >= 0)')\n",
+    "assert by_line[4] == '-4' and by_line[3] == '-3'\n",
+    "assert (-7) // 2 == -4 and 7 // -2 == -4            # python floors: differs from Lean on 7 / (-2)\n",
+    "\n",
+    "print('--- autoImplicit: a typo silently becomes a new universally-quantified variable ---')\n",
+    "for m in lean_msgs(\"theorem t_typo (x : ℕ) : x + 1 > xx := by sorry\\n#check @t_typo\")[0]: print(m['severity'], m['data'][:300])\n",
+    "print('--- the same with autoImplicit off: an error, as it should be ---')\n",
+    "for m in lean_msgs(\"set_option autoImplicit false in\\ntheorem t_typo2 (x : ℕ) : x + 1 > xx := by sorry\")[0]: print(m['severity'], m['data'][:200])\n",
+    "L.close()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

code/notebooks/0303-hoare-refinement.ipynb

@@ -0,0 +1,852 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "8fc053db",
+   "metadata": {},
+   "source": [
+    "# 0303 - Hoare logic and refinement; verifying a small function in Lean 4\n",
+    "\n",
+    "Scratch-pad for note `0303 - Hoare Logic.pdf`. Naive first, efficient second, assert agreement.\n",
+    "\n",
+    "**A. A weakest-precondition calculus for a tiny `while` language** (sympy for substitution).\n",
+    "* NAIVE = *run* the program on every state of a box and check the postcondition (semantic weakest precondition, `wp_sem`).\n",
+    "* EFFICIENT = *compute* `wp` by the syntactic rules (substitution) and verification conditions (VCs) for loops.\n",
+    "* Assert: `wp_syn == wp_sem` on 300 random loop-free programs; VCs valid => the Hoare triple holds on every box state;\n",
+    "  wrong invariants are rejected, and we classify the counterexample (reachable = real bug, unreachable = proof artefact).\n",
+    "\n",
+    "**B. A verified linear search in Lean 4** (the DafnyBench Fig. 2 `LinearSearch`): code, loop invariant as an induction\n",
+    "hypothesis, theorem, and what Lean says about an off-by-one variant. **C. Integers**: spec over Z, code over int32.\n",
+    "\n",
+    "No GPU. Lean: ONE worker via `align.verify.lean`, nice 15, after the MemAvailable>=45 GB check in the helper."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "7964bb54",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:55:15.064599Z",
+     "iopub.status.busy": "2026-10-14T09:55:15.064441Z",
+     "iopub.status.idle": "2026-10-14T09:55:15.330291Z",
+     "shell.execute_reply": "2026-10-14T09:55:15.329450Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
+    "import itertools, random, sys, time\n",
+    "import sympy as sp\n",
+    "random.seed(0)\n",
+    "n, i, s, x, y = sp.symbols('n i s x y', integer=True)\n",
+    "\n",
+    "# AST:  ('skip',) | ('assign', var, expr) | ('seq', S1, S2) | ('if', cond, S1, S2) | ('while', cond, inv, variant, body)\n",
+    "def seq(*ss):\n",
+    "    out = ss[-1]\n",
+    "    for t in reversed(ss[:-1]): out = ('seq', t, out)\n",
+    "    return out\n",
+    "def imp(a, b): return sp.Or(sp.Not(a), b)\n",
+    "\n",
+    "_cache = {}\n",
+    "def ev(expr, st):\n",
+    "    key = (expr, tuple(sorted(st, key=str)))\n",
+    "    if key not in _cache:\n",
+    "        syms = sorted(st, key=str)\n",
+    "        _cache[key] = (syms, sp.lambdify(syms, expr, 'math'))\n",
+    "    syms, f = _cache[key]\n",
+    "    return f(*[st[v] for v in syms])\n",
+    "\n",
+    "class NonTerminating(Exception): pass\n",
+    "\n",
+    "def run(S, st, fuel=10_000):\n",
+    "    \"\"\"Big-step semantics <S, st> => st'. Returns the final state; raises NonTerminating when fuel is out.\"\"\"\n",
+    "    st = dict(st)\n",
+    "    def go(S):\n",
+    "        nonlocal fuel\n",
+    "        k = S[0]\n",
+    "        if k == 'skip': return\n",
+    "        if k == 'assign': st[S[1]] = ev(S[2], st); return\n",
+    "        if k == 'seq': go(S[1]); go(S[2]); return\n",
+    "        if k == 'if': go(S[2] if ev(S[1], st) else S[3]); return\n",
+    "        if k == 'while':\n",
+    "            while ev(S[1], st):\n",
+    "                fuel -= 1\n",
+    "                if fuel < 0: raise NonTerminating\n",
+    "                go(S[4])\n",
+    "    go(S)\n",
+    "    return st\n",
+    "\n",
+    "def wp(S, Q):\n",
+    "    \"\"\"Weakest precondition (loop-free fragment, plus `while` handled via its invariant: wp = I, see vcs()).\"\"\"\n",
+    "    k = S[0]\n",
+    "    if k == 'skip': return Q\n",
+    "    if k == 'assign': return Q.subs(S[1], S[2], simultaneous=True)\n",
+    "    if k == 'seq': return wp(S[1], wp(S[2], Q))\n",
+    "    if k == 'if': return sp.Or(sp.And(S[1], wp(S[2], Q)), sp.And(sp.Not(S[1]), wp(S[3], Q)))\n",
+    "    if k == 'while': return S[2]                       # annotated invariant; its VCs are generated by vcs()\n",
+    "    raise ValueError(k)\n",
+    "\n",
+    "_ghost_ctr = itertools.count()\n",
+    "\n",
+    "def vcs(S, Q):\n",
+    "    \"\"\"Verification conditions that make wp(S,Q) (with loops replaced by their invariants) sound. Returns [(name, formula)].\"\"\"\n",
+    "    k = S[0]\n",
+    "    if k in ('skip', 'assign'): return []\n",
+    "    if k == 'seq':  return vcs(S[1], wp(S[2], Q)) + vcs(S[2], Q)\n",
+    "    if k == 'if':   return vcs(S[2], Q) + vcs(S[3], Q)\n",
+    "    if k == 'while':\n",
+    "        b, I, V, body = S[1], S[2], S[3], S[4]\n",
+    "        v0 = S[5] if len(S) > 5 else sp.Symbol(f'v0_{next(_ghost_ctr)}', integer=True)     # one ghost per loop (nested loops must not share it)\n",
+    "        ante = sp.And(I, b) if V is None else sp.And(I, b, sp.Eq(V, v0))     # v0 = the (ghost) value of the variant at loop entry\n",
+    "        out = [('preserve', imp(ante, wp(body, I))),                    # I is preserved (the I-part of the body VC, reported separately)\n",
+    "               ('exit',     imp(sp.And(I, sp.Not(b)), Q))]\n",
+    "        if V is not None:\n",
+    "            R = sp.And(I, V < v0)                                       # ONE postcondition for the body: I and the variant drop\n",
+    "            out += [('variant>=0', imp(I, V >= 0)),\n",
+    "                    ('variant-decreases', imp(sp.And(I, b, sp.Eq(V, v0)), wp(body, R)))]\n",
+    "            return out + vcs(body, R)                                   # (needed for nested loops: their exit VCs must imply V < v0)\n",
+    "        return out + vcs(body, I)\n",
+    "    raise ValueError(k)\n",
+    "\n",
+    "def box(vars_, lo, hi):\n",
+    "    return [dict(zip(vars_, p)) for p in itertools.product(range(lo, hi + 1), repeat=len(vars_))]\n",
+    "\n",
+    "GHOST = range(-12, 13)\n",
+    "def holds(F, st, ghost=None):\n",
+    "    \"\"\"Validity of F at the state st. Free symbols that are not program variables (the ghost values v0 of the variant) are\n",
+    "    universally quantified over a range. (Before the fix they were silently left unbound and Eq(V, v0) was structurally False,\n",
+    "    which made every 'variant-decreases' VC vacuously true.)\"\"\"\n",
+    "    ghosts = sorted(F.free_symbols - set(st), key=str)\n",
+    "    if not ghosts: return bool(ev(F, st))\n",
+    "    return all(bool(ev(F, {**st, **dict(zip(ghosts, gv))})) for gv in itertools.product(ghost or GHOST, repeat=len(ghosts)))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c8f54cc4",
+   "metadata": {},
+   "source": [
+    "## A1. Straight-line and conditional code: `wp` by substitution equals the semantic weakest precondition\n",
+    "Semantic wp: `{ st | S terminates from st and the final state satisfies Q }`. We compare the *set* of box states on both sides."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "fe2daf34",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:55:15.332288Z",
+     "iopub.status.busy": "2026-10-14T09:55:15.332056Z",
+     "iopub.status.idle": "2026-10-14T09:55:19.396869Z",
+     "shell.execute_reply": "2026-10-14T09:55:19.395883Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "300 random loop-free programs x 49 states = 14700 (program,state) pairs; wp_syn != wp_sem on 0   (4.0s)\n"
+     ]
+    }
+   ],
+   "source": [
+    "VARS = [x, y]\n",
+    "EXPRS = [x + 1, x - 1, 2 * x, y, x + y, x - y, y + 3, sp.Integer(0), x * y]\n",
+    "CONDS = [sp.Lt(x, y), sp.Le(x, 0), sp.Eq(x, y), sp.Gt(y, 1), sp.Ne(x + y, 2)]\n",
+    "POSTS = [sp.Eq(x, y), sp.Ge(x, 0), sp.Lt(x + y, 4), sp.Eq(2 * x, y + 1), sp.And(sp.Ge(x, 0), sp.Ge(y, 0)), sp.Ne(x, y)]\n",
+    "\n",
+    "def rand_prog(d):\n",
+    "    r = random.random()\n",
+    "    if d == 0 or r < 0.35: return ('assign', random.choice(VARS), random.choice(EXPRS))\n",
+    "    if r < 0.7: return ('seq', rand_prog(d - 1), rand_prog(d - 1))\n",
+    "    return ('if', random.choice(CONDS), rand_prog(d - 1), rand_prog(d - 1))\n",
+    "\n",
+    "STATES = box(VARS, -3, 3)\n",
+    "checked = mism = 0\n",
+    "t0 = time.time()\n",
+    "for _ in range(300):\n",
+    "    S = rand_prog(3); Q = random.choice(POSTS)\n",
+    "    W = wp(S, Q)\n",
+    "    for st in STATES:\n",
+    "        syn = holds(W, st)\n",
+    "        sem = holds(Q, run(S, st))               # loop-free: always terminates\n",
+    "        checked += 1; mism += (syn != sem)\n",
+    "assert mism == 0\n",
+    "print(f'300 random loop-free programs x {len(STATES)} states = {checked} (program,state) pairs; wp_syn != wp_sem on {mism}   ({time.time()-t0:.1f}s)')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c520492b",
+   "metadata": {},
+   "source": [
+    "## A2. Loops: invariant + variant. The worked example `sum`\n",
+    "```\n",
+    "{n >= 0}  i := 0; s := 0; while i < n  [inv I, variant n - i]  do  s := s + i; i := i + 1  {2 s = n (n - 1)}\n",
+    "```\n",
+    "`I : 0 <= i <= n  and  2 s = i (i - 1)`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "cb9b378a",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:55:19.398693Z",
+     "iopub.status.busy": "2026-10-14T09:55:19.398538Z",
+     "iopub.status.idle": "2026-10-14T09:55:21.057597Z",
+     "shell.execute_reply": "2026-10-14T09:55:21.056648Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "correct invariant                  all VCs valid on box: True\n",
+      "semantic run on 48 inputs satisfying/not P: triple holds -> True\n"
+     ]
+    }
+   ],
+   "source": [
+    "body = seq(('assign', s, s + i), ('assign', i, i + 1))\n",
+    "cond = sp.Lt(i, n)\n",
+    "I_good = sp.And(sp.Ge(i, 0), sp.Le(i, n), sp.Eq(2 * s, i * (i - 1)))\n",
+    "P = sp.Ge(n, 0)\n",
+    "Qpost = sp.Eq(2 * s, n * (n - 1))\n",
+    "def program(inv, variant=n - i, bodyS=body):\n",
+    "    return seq(('assign', i, sp.Integer(0)), ('assign', s, sp.Integer(0)), ('while', cond, inv, variant, bodyS))\n",
+    "\n",
+    "BOX = box([n, i, s], -2, 9)\n",
+    "BOX = [st for st in BOX if -2 <= st[s] <= 40] if False else BOX\n",
+    "\n",
+    "def verify(S, P, Q, label, reachable_states=None):\n",
+    "    \"\"\"Generate all VCs (including P => wp(S,Q)) and discharge each on the box. Return dict name->counterexample or None.\"\"\"\n",
+    "    conds = [('pre=>wp', imp(P, wp(S, Q)))] + vcs(S, Q)\n",
+    "    res = {}\n",
+    "    for name, F in conds:\n",
+    "        bad = next((st for st in BOX if not holds(F, st)), None)\n",
+    "        res[name] = bad\n",
+    "    ok = all(v is None for v in res.values())\n",
+    "    print(f'{label:34s} all VCs valid on box: {ok}' + ('' if ok else '   failing: ' + ', '.join(f'{k}@{ {str(a):b for a,b in v.items()} }' for k, v in res.items() if v)))\n",
+    "    return res\n",
+    "\n",
+    "def triple_holds_semantically(S, P, Q, inputs):\n",
+    "    \"\"\"Naive check: run from every input state satisfying P; the final state must satisfy Q (and the run must terminate).\"\"\"\n",
+    "    for st in inputs:\n",
+    "        if not holds(P, st): continue\n",
+    "        try: fin = run(S, st)\n",
+    "        except NonTerminating: return False, st\n",
+    "        if not holds(Q, fin): return False, st\n",
+    "    return True, None\n",
+    "\n",
+    "INPUTS = [dict(n=a, i=b, s=c) for a in range(-2, 10) for b in (0, 5) for c in (0, 7)]\n",
+    "INPUTS = [{n: d['n'], i: d['i'], s: d['s']} for d in INPUTS]\n",
+    "\n",
+    "r_good = verify(program(I_good), P, Qpost, 'correct invariant')\n",
+    "ok_sem, ce = triple_holds_semantically(program(I_good), P, Qpost, INPUTS)\n",
+    "assert all(v is None for v in r_good.values()) and ok_sem\n",
+    "print('semantic run on', len(INPUTS), 'inputs satisfying/not P: triple holds ->', ok_sem)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1b414665",
+   "metadata": {},
+   "source": [
+    "### A3. Five ways to fail, and what each failure means\n",
+    "1. **Weak invariant** (`0 <= i <= n` only): the *exit* VC fails. The program is correct; the proof attempt fails. The counterexample state is\n",
+    "   *unreachable* from any state satisfying `P`.\n",
+    "2. **Too-strong invariant** (`i < n`): `pre=>wp` / `preserve` fail. Program correct, annotation wrong.\n",
+    "3. **Buggy program** (`s := s + i + 1`): VCs fail with the *right* invariant; running the program from a concrete input violates the post. A **real bug**.\n",
+    "4. **Wrong postcondition** (`2 s = n (n + 1)`): the spec is wrong for the code; both checks fail."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "a6bbfb08",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:55:21.059329Z",
+     "iopub.status.busy": "2026-10-14T09:55:21.059160Z",
+     "iopub.status.idle": "2026-10-14T09:55:26.415346Z",
+     "shell.execute_reply": "2026-10-14T09:55:26.414446Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1 weak invariant                   all VCs valid on box: False   failing: exit@{'n': 0, 'i': 0, 's': -2}\n",
+      "   counterexample {'n': 0, 'i': 0, 's': -2} -> reachable by the real program? False\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2 too-strong invariant             all VCs valid on box: False   failing: pre=>wp@{'n': 0, 'i': -2, 's': -2}, preserve@{'n': 1, 'i': 0, 's': 0}, variant-decreases@{'n': 1, 'i': 0, 's': 0}\n",
+      "   but the program is correct (semantic check): True\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "3 buggy program, right invariant   all VCs valid on box: False   failing: preserve@{'n': 1, 'i': 0, 's': 0}, variant-decreases@{'n': 1, 'i': 0, 's': 0}\n",
+      "   concrete failing input from the run: {'n': 1, 'i': 0, 's': 0} -> confirmed bug\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "4 wrong postcondition              all VCs valid on box: False   failing: exit@{'n': 1, 'i': 1, 's': 0}\n",
+      "   run also violates it, e.g. {'n': 1, 'i': 0, 's': 0}\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "5 wrong variant (V = i, increasing) all VCs valid on box: False   failing: variant-decreases@{'n': 1, 'i': 0, 's': 0}\n",
+      "   the program is still correct and terminates (semantic check): True -> the failing VC is about the PROOF of termination, not about a bug\n"
+     ]
+    }
+   ],
+   "source": [
+    "reachable = set()\n",
+    "for nv in range(0, 10):\n",
+    "    st = {n: nv, i: 0, s: 0}\n",
+    "    # collect every loop-head state of the correct program\n",
+    "    cur = {n: nv, i: 0, s: 0}\n",
+    "    while cur[i] < cur[n]:\n",
+    "        reachable.add((cur[n], cur[i], cur[s])); cur[s] += cur[i]; cur[i] += 1\n",
+    "    reachable.add((cur[n], cur[i], cur[s]))\n",
+    "\n",
+    "I_weak = sp.And(sp.Ge(i, 0), sp.Le(i, n))\n",
+    "r1 = verify(program(I_weak), P, Qpost, '1 weak invariant')\n",
+    "assert r1['exit'] is not None and all(r1[k] is None for k in ('pre=>wp', 'preserve'))\n",
+    "ce1 = r1['exit']; print('   counterexample', {str(k): v for k, v in ce1.items()}, '-> reachable by the real program?', (ce1[n], ce1[i], ce1[s]) in reachable)\n",
+    "assert (ce1[n], ce1[i], ce1[s]) not in reachable           # a proof artefact, NOT a bug\n",
+    "\n",
+    "I_strong = sp.And(sp.Ge(i, 0), sp.Lt(i, n), sp.Eq(2 * s, i * (i - 1)))\n",
+    "r2 = verify(program(I_strong), P, Qpost, '2 too-strong invariant')\n",
+    "assert any(v is not None for v in r2.values())\n",
+    "print('   but the program is correct (semantic check):', triple_holds_semantically(program(I_good), P, Qpost, INPUTS)[0])\n",
+    "\n",
+    "bug_body = seq(('assign', s, s + i + 1), ('assign', i, i + 1))\n",
+    "r3 = verify(program(I_good, bodyS=bug_body), P, Qpost, '3 buggy program, right invariant')\n",
+    "ok3, ce3 = triple_holds_semantically(program(I_good, bodyS=bug_body), P, Qpost, INPUTS)\n",
+    "assert r3['preserve'] is not None and not ok3\n",
+    "print('   concrete failing input from the run:', {str(k): v for k, v in ce3.items()}, '-> confirmed bug')\n",
+    "\n",
+    "Q_wrong = sp.Eq(2 * s, n * (n + 1))\n",
+    "r4 = verify(program(I_good), P, Q_wrong, '4 wrong postcondition')\n",
+    "ok4, ce4 = triple_holds_semantically(program(I_good), P, Q_wrong, INPUTS)\n",
+    "assert r4['exit'] is not None and not ok4\n",
+    "print('   run also violates it, e.g.', {str(k): v for k, v in ce4.items()})\n",
+    "\n",
+    "# 5. a WRONG VARIANT (V = i grows): the invariant and the code are right; only the termination argument is wrong\n",
+    "r5 = verify(program(I_good, variant=i), P, Qpost, '5 wrong variant (V = i, increasing)')\n",
+    "assert r5['variant-decreases'] is not None and r5['preserve'] is None and r5['exit'] is None\n",
+    "print('   the program is still correct and terminates (semantic check):', triple_holds_semantically(program(I_good), P, Qpost, INPUTS)[0],\n",
+    "      '-> the failing VC is about the PROOF of termination, not about a bug')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "df54905f",
+   "metadata": {},
+   "source": [
+    "### A4. Soundness of the `while` rule, checked numerically\n",
+    "For every loop program whose VCs (all four) are valid on the box, every box state satisfying the invariant and variant terminates in `Q`\n",
+    "(this is `I  subseteq  wp(while, Q)`, the total-correctness statement of Boxed result 3)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "cd82f6da",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:55:26.416985Z",
+     "iopub.status.busy": "2026-10-14T09:55:26.416825Z",
+     "iopub.status.idle": "2026-10-14T09:55:26.469647Z",
+     "shell.execute_reply": "2026-10-14T09:55:26.468775Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "40 box states satisfying I and variant>=0: every one terminates in Q (0 violations)\n"
+     ]
+    }
+   ],
+   "source": [
+    "cnt = viol = 0\n",
+    "for st in BOX:\n",
+    "    if holds(I_good, st) and holds(n - i >= 0, st):\n",
+    "        S_loop = ('while', cond, I_good, n - i, body)\n",
+    "        fin = run(S_loop, st)                          # raises if it does not terminate\n",
+    "        cnt += 1; viol += (not holds(Qpost, fin))\n",
+    "assert viol == 0\n",
+    "print(f'{cnt} box states satisfying I and variant>=0: every one terminates in Q ({viol} violations)')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c5b343fb",
+   "metadata": {},
+   "source": [
+    "### A5. Nested loops (why the body is verified against `I and V < v0` in ONE postcondition, with one ghost per loop)\n",
+    "```\n",
+    "{n >= 0 and m >= 0}  i := 0; s := 0; while i < n [I1, V1 = n-i] do ( j := 0; while j < m [I2, V2 = m-j] do ( s := s+1; j := j+1 ); i := i+1 )  {s = n m}\n",
+    "I1 = 0<=i<=n and m>=0 and s = i m          I2 = 0<=j<=m and 0<=i<n and m>=0 and s = i m + j and n-i = g1\n",
+    "```\n",
+    "`g1` is the ghost (`v0`) of the outer variant: the INNER loop's exit VC must imply `V1 < g1` for the outer body, so the inner invariant has to carry `n - i = g1`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "b71a2c7b",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:55:26.471193Z",
+     "iopub.status.busy": "2026-10-14T09:55:26.471039Z",
+     "iopub.status.idle": "2026-10-14T09:55:35.643578Z",
+     "shell.execute_reply": "2026-10-14T09:55:35.642581Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "9 VCs generated for the nested program: ['pre=>wp', 'preserve', 'exit', 'variant>=0', 'variant-decreases', 'preserve', 'exit', 'variant>=0', 'variant-decreases']\n",
+      "all VCs valid on the box; the program run from 9 inputs terminates in Q (s = n m)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "wrong inner variant (V2 = j): failing VCs -> ['variant-decreases']\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "m_, j_, g1, g2 = sp.symbols('m j g1 g2', integer=True)\n",
+    "cond1, cond2 = sp.Lt(i, n), sp.Lt(j_, m_)\n",
+    "I1 = sp.And(sp.Ge(i, 0), sp.Le(i, n), sp.Ge(m_, 0), sp.Eq(s, i * m_))\n",
+    "I2 = sp.And(sp.Ge(j_, 0), sp.Le(j_, m_), sp.Ge(i, 0), sp.Lt(i, n), sp.Ge(m_, 0), sp.Eq(s, i * m_ + j_), sp.Eq(n - i, g1))\n",
+    "def nested(inner_variant):\n",
+    "    inner = ('while', cond2, I2, inner_variant, seq(('assign', s, s + 1), ('assign', j_, j_ + 1)), g2)\n",
+    "    outer_body = seq(('assign', j_, sp.Integer(0)), inner, ('assign', i, i + 1))\n",
+    "    return seq(('assign', i, sp.Integer(0)), ('assign', s, sp.Integer(0)), ('while', cond1, I1, n - i, outer_body, g1))\n",
+    "P_n, Q_n = sp.And(sp.Ge(n, 0), sp.Ge(m_, 0)), sp.Eq(s, n * m_)\n",
+    "BOX_N = [dict(zip([n, m_, i, j_, s], p)) for p in itertools.product(range(0, 3), range(0, 3), range(0, 4), range(0, 4), range(0, 10))]\n",
+    "GN = range(-3, 4)\n",
+    "def failing_vcs(S):\n",
+    "    conds = [('pre=>wp', imp(P_n, wp(S, Q_n)))] + vcs(S, Q_n)\n",
+    "    return [name for name, F in conds if any(not holds(F, st, GN) for st in BOX_N)], [name for name, _ in conds]\n",
+    "bad, names = failing_vcs(nested(m_ - j_))\n",
+    "print(len(names), 'VCs generated for the nested program:', names)\n",
+    "assert bad == [], bad\n",
+    "ins = [st for st in BOX_N if holds(P_n, st) and st[i] == 0 and st[j_] == 0 and st[s] == 0]\n",
+    "for st in ins:\n",
+    "    fin = run(nested(m_ - j_), st); assert holds(Q_n, fin)           # raises NonTerminating if the run does not terminate\n",
+    "print(f'all VCs valid on the box; the program run from {len(ins)} inputs terminates in Q (s = n m)')\n",
+    "bad2, _ = failing_vcs(nested(j_))                                      # a WRONG inner variant (V2 = j grows)\n",
+    "print('wrong inner variant (V2 = j): failing VCs ->', bad2)\n",
+    "assert 'variant-decreases' in bad2\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "63e05762",
+   "metadata": {},
+   "source": [
+    "## B. Verifying `LinearSearch` in Lean 4 (DafnyBench Fig. 2 in Lean)\n",
+    "Dafny: `ensures 0 <= n <= a.Length; ensures n == a.Length || P(a[n]); ensures forall i :: 0 <= i < n ==> !P(a[i])` with loop invariants\n",
+    "`0 <= n <= a.Length` and `forall i :: 0 <= i < n ==> !P(a[i])`. Abstracting `a[i]` away (we search an index predicate `P`):"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "55dab72a",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:55:35.645474Z",
+     "iopub.status.busy": "2026-10-14T09:55:35.645313Z",
+     "iopub.status.idle": "2026-10-14T09:55:40.310783Z",
+     "shell.execute_reply": "2026-10-14T09:55:40.309894Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[info] line 36: 'search_correct' depends on axioms: [propext, Quot.sound]\n",
+      "[info] line 37: 3\n",
+      "[info] line 38: 10\n"
+     ]
+    }
+   ],
+   "source": [
+    "sys.path.insert(0, os.path.abspath('.'))\n",
+    "from _lean_helper import Lean\n",
+    "L = Lean()\n",
+    "code = r'''\n",
+    "/-- the loop `while n < len do if P n then return else n := n + 1`, as a tail-recursive function -/\n",
+    "def search (P : Nat → Bool) (len : Nat) (n : Nat) : Nat :=\n",
+    "  if n < len then (if P n then n else search P len (n + 1)) else n\n",
+    "termination_by len - n\n",
+    "\n",
+    "/-- postcondition of Dafny's LinearSearch (with the start index `n0` of the loop) -/\n",
+    "def Post (P : Nat → Bool) (len : Nat) (r : Nat) : Prop :=\n",
+    "  r ≤ len ∧ (r = len ∨ P r = true) ∧ ∀ i, i < r → P i = false\n",
+    "\n",
+    "/-- loop INVARIANT as an induction hypothesis: from any state satisfying (n ≤ len ∧ ∀ i < n, ¬P i) the loop ends in Post -/\n",
+    "theorem search_inv (P : Nat → Bool) (len : Nat) :\n",
+    "    ∀ n, n ≤ len → (∀ i, i < n → P i = false) → Post P len (search P len n) := by\n",
+    "  intro n\n",
+    "  induction' h : len - n using Nat.strong_induction_on with k ih generalizing n\n",
+    "  intro hn hinv\n",
+    "  unfold search\n",
+    "  by_cases hlt : n < len\n",
+    "  · rw [if_pos hlt]\n",
+    "    by_cases hp : P n = true\n",
+    "    · rw [if_pos hp]\n",
+    "      exact ⟨by omega, Or.inr hp, hinv⟩\n",
+    "    · rw [if_neg hp]\n",
+    "      have hp' : P n = false := by simpa using hp\n",
+    "      apply ih (len - (n + 1)) (by omega) (n + 1) rfl (by omega)\n",
+    "      intro j hj\n",
+    "      by_cases hjn : j = n\n",
+    "      · subst hjn; exact hp'\n",
+    "      · exact hinv j (by omega)\n",
+    "  · rw [if_neg hlt]\n",
+    "    exact ⟨by omega, Or.inl (by omega), hinv⟩\n",
+    "\n",
+    "theorem search_correct (P : Nat → Bool) (len : Nat) : Post P len (search P len 0) :=\n",
+    "  search_inv P len 0 (Nat.zero_le _) (by intro i hi; omega)\n",
+    "\n",
+    "#print axioms search_correct\n",
+    "#eval search (fun i => i == 3) 10 0\n",
+    "#eval search (fun i => i == 30) 10 0\n",
+    "'''\n",
+    "r = L.show(code, timeout=300, keep=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "041614a3",
+   "metadata": {},
+   "source": [
+    "### B2. The same proof attempt on an off-by-one loop (`n + 1 < len`)\n",
+    "The proof script is *identical*; Lean refuses. A failed proof is not yet a bug report (note 0304) -- but here a one-line `#eval` confirms it."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "3f1bb803",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:55:40.312728Z",
+     "iopub.status.busy": "2026-10-14T09:55:40.312554Z",
+     "iopub.status.idle": "2026-10-14T09:55:40.878639Z",
+     "shell.execute_reply": "2026-10-14T09:55:40.877856Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[ok] no messages\n",
+      "[error] line 21: omega could not prove the goal:\n",
+      "a possible counterexample may satisfy the constraints\n",
+      "  a ≥ 0\n",
+      "where\n",
+      " a := ↑n\n",
+      "--- concrete confirmation (no match anywhere: the correct loop returns len = 10) ---\n",
+      "[info] line 2: 10\n",
+      "[info] line 3: 9\n",
+      "[info] line 4: false\n",
+      "[info] line 5: true\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "{'messages': [{'severity': 'info',\n",
+       "   'pos': {'line': 2, 'column': 0},\n",
+       "   'endPos': {'line': 2, 'column': 5},\n",
+       "   'data': '10'},\n",
+       "  {'severity': 'info',\n",
+       "   'pos': {'line': 3, 'column': 0},\n",
+       "   'endPos': {'line': 3, 'column': 5},\n",
+       "   'data': '9'},\n",
+       "  {'severity': 'info',\n",
+       "   'pos': {'line': 4, 'column': 0},\n",
+       "   'endPos': {'line': 4, 'column': 5},\n",
+       "   'data': 'false'},\n",
+       "  {'severity': 'info',\n",
+       "   'pos': {'line': 5, 'column': 0},\n",
+       "   'endPos': {'line': 5, 'column': 5},\n",
+       "   'data': 'true'}],\n",
+       " 'env': 5}"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "L.show(r'''\n",
+    "def searchBad (P : Nat → Bool) (len : Nat) (n : Nat) : Nat :=\n",
+    "  if n + 1 < len then (if P n then n else searchBad P len (n + 1)) else n\n",
+    "termination_by len - n\n",
+    "instance (P : Nat → Bool) (len r : Nat) : Decidable (Post P len r) := by unfold Post; infer_instance\n",
+    "''', keep=True)\n",
+    "bad = r'''\n",
+    "theorem searchBad_inv (P : Nat → Bool) (len : Nat) :\n",
+    "    ∀ n, n ≤ len → (∀ i, i < n → P i = false) → Post P len (searchBad P len n) := by\n",
+    "  intro n\n",
+    "  induction' h : len - n using Nat.strong_induction_on with k ih generalizing n\n",
+    "  intro hn hinv\n",
+    "  unfold searchBad\n",
+    "  by_cases hlt : n + 1 < len\n",
+    "  · rw [if_pos hlt]\n",
+    "    by_cases hp : P n = true\n",
+    "    · rw [if_pos hp]\n",
+    "      exact ⟨by omega, Or.inr hp, hinv⟩\n",
+    "    · rw [if_neg hp]\n",
+    "      have hp' : P n = false := by simpa using hp\n",
+    "      apply ih (len - (n + 1)) (by omega) (n + 1) rfl (by omega)\n",
+    "      intro j hj\n",
+    "      by_cases hjn : j = n\n",
+    "      · subst hjn; exact hp'\n",
+    "      · exact hinv j (by omega)\n",
+    "  · rw [if_neg hlt]\n",
+    "    exact ⟨by omega, Or.inl (by omega), hinv⟩\n",
+    "'''\n",
+    "rr = L.show(bad, timeout=300)\n",
+    "print('--- concrete confirmation (no match anywhere: the correct loop returns len = 10) ---')\n",
+    "L.show('''\n",
+    "#eval search (fun _ => false) 10 0\n",
+    "#eval searchBad (fun _ => false) 10 0\n",
+    "#eval decide (Post (fun _ => false) 10 (searchBad (fun _ => false) 10 0))\n",
+    "#eval decide (Post (fun _ => false) 10 (search (fun _ => false) 10 0))\n",
+    "''')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "acbe59cb",
+   "metadata": {},
+   "source": [
+    "## C. The spec lives in Z, the code in int32: binary-search midpoint\n",
+    "VC (over Z): `lo <= hi  =>  lo <= (lo+hi)/2 <= hi`. Over 8-bit words the identical VC is **false** exactly when `lo + hi >= 256`.\n",
+    "Derivation: wrapped `lo+hi = lo+hi-256`, so `mid = (lo+hi-256)/2 < lo  <=>  hi - lo < 256`, always true for 8-bit words.\n",
+    "Count the violating pairs in closed form, by brute force in Python, and by Lean `#eval` on `UInt8`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "86a322fd",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:55:40.880806Z",
+     "iopub.status.busy": "2026-10-14T09:55:40.880641Z",
+     "iopub.status.idle": "2026-10-14T09:55:41.086181Z",
+     "shell.execute_reply": "2026-10-14T09:55:41.085341Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "UInt8 pairs lo<=hi: 32896;  violating VC: unsafe 16384 (= #{lo+hi>=256} = 16384),  safe 0\n",
+      "Lean #eval counts (unsafe, safe, UInt32 mid): ['16384', '0', '852516352'] ; decide-proof of the 32-bit violation:  True\n",
+      "int32: lo = 1073741829  hi = 2147483646  (lo+hi) wraps to -1073741821 -> mid = -536870911 which is < lo: True\n"
+     ]
+    }
+   ],
+   "source": [
+    "M = 256\n",
+    "def mid_unsafe(lo, hi): return ((lo + hi) % M) // 2                    # what UInt8 arithmetic computes\n",
+    "def mid_safe(lo, hi):   return (lo + (hi - lo) // 2) % M              # no intermediate overflow when lo <= hi\n",
+    "pairs = [(lo, hi) for lo in range(M) for hi in range(lo, M)]\n",
+    "py_bad_unsafe = sum(1 for lo, hi in pairs if not (lo <= mid_unsafe(lo, hi) <= hi))\n",
+    "py_bad_safe   = sum(1 for lo, hi in pairs if not (lo <= mid_safe(lo, hi) <= hi))\n",
+    "closed_form   = sum(1 for lo, hi in pairs if lo + hi >= M)             # the derived characterisation\n",
+    "assert py_bad_unsafe == closed_form and py_bad_safe == 0\n",
+    "print(f'UInt8 pairs lo<=hi: {len(pairs)};  violating VC: unsafe {py_bad_unsafe} (= #{{lo+hi>=256}} = {closed_form}),  safe {py_bad_safe}')\n",
+    "# the same count in Lean on UInt8\n",
+    "msgs = L.run('''\n",
+    "#eval ((List.range 256).foldl (fun acc lo => (List.range 256).foldl (fun acc hi =>\n",
+    "    if lo ≤ hi ∧ ¬ (UInt8.ofNat lo ≤ (UInt8.ofNat lo + UInt8.ofNat hi) / 2 ∧ (UInt8.ofNat lo + UInt8.ofNat hi) / 2 ≤ UInt8.ofNat hi) then acc + 1 else acc) acc) 0)\n",
+    "#eval ((List.range 256).foldl (fun acc lo => (List.range 256).foldl (fun acc hi =>\n",
+    "    if lo ≤ hi ∧ ¬ (UInt8.ofNat lo ≤ UInt8.ofNat lo + (UInt8.ofNat hi - UInt8.ofNat lo) / 2 ∧ UInt8.ofNat lo + (UInt8.ofNat hi - UInt8.ofNat lo) / 2 ≤ UInt8.ofNat hi) then acc + 1 else acc) acc) 0)\n",
+    "#eval ((3000000000 : UInt32) + 3000000000) / 2\n",
+    "example : ¬ ((3000000000 : UInt32) ≤ (3000000000 + 3000000000) / 2) := by decide\n",
+    "''')['messages']\n",
+    "lean_counts = [m['data'].strip() for m in msgs if m['severity'] == 'info']\n",
+    "print('Lean #eval counts (unsafe, safe, UInt32 mid):', lean_counts, '; decide-proof of the 32-bit violation: ', [m['severity'] for m in msgs if m['severity'] == 'error'] == [])\n",
+    "assert int(lean_counts[0]) == py_bad_unsafe and int(lean_counts[1]) == 0\n",
+    "# 32-bit: python model of int32 wrap, the famous boundary case\n",
+    "def wrap32(v): return (v + 2**31) % 2**32 - 2**31\n",
+    "lo, hi = 2**30 + 5, 2**31 - 2\n",
+    "print('int32: lo =', lo, ' hi =', hi, ' (lo+hi) wraps to', wrap32(lo + hi), '-> mid =', wrap32(lo + hi) // 2, 'which is < lo:', wrap32(lo + hi) // 2 < lo)\n",
+    "assert wrap32(lo + hi) // 2 < lo and lo <= (lo + hi) // 2 <= hi"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f9bbfb94",
+   "metadata": {},
+   "source": [
+    "## D. Refinement: `x :[pre, post]  is-refined-by  x := e   iff   pre => post[e/x]`\n",
+    "Predicate-transformer reading: `wp_spec(Q)(s) = pre(s) and forall x'. post(s, x') => Q(x')`, `wp_assign(Q)(s) = Q(e(s))`;\n",
+    "`spec ⊑ impl` iff `wp_spec(Q) ⊆ wp_impl(Q)` for every `Q`. We decide this by brute force over ALL 2^11 postconditions on a box\n",
+    "(named approximation: box of 11 values) and compare with the one-line law."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "ca6a9e0e",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:55:41.087953Z",
+     "iopub.status.busy": "2026-10-14T09:55:41.087795Z",
+     "iopub.status.idle": "2026-10-14T09:55:42.520867Z",
+     "shell.execute_reply": "2026-10-14T09:55:42.519895Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "80 random (spec, assignment) pairs: all-Q brute force == worst-Q shortcut == law pre => post[e/x]; refined in 71, not refined in 9\n"
+     ]
+    }
+   ],
+   "source": [
+    "BOXX = list(range(-5, 6))                   # final-state values x' (11 of them => 2^11 postconditions Q)\n",
+    "START = list(range(-2, 3))                  # start states s0; |b|<=2, |k|<=1 keep e(s0) inside BOXX\n",
+    "def refines_bruteforce(pre, post, e):\n",
+    "    # for every Q (every subset of the box), wp_spec(Q) ⊆ wp_assign(Q) on every start state\n",
+    "    for s0 in START:\n",
+    "        if not pre(s0): continue                                  # wp_spec(Q)(s0)=False for all Q: nothing to check\n",
+    "        allowed = [xp for xp in BOXX if post(s0, xp)]             # outcomes the spec permits\n",
+    "        # wp_spec(Q)(s0) holds iff allowed ⊆ Q ; refinement needs e(s0) ∈ Q for all such Q.  Worst Q = allowed.\n",
+    "        if e(s0) not in allowed: return False\n",
+    "    return True\n",
+    "def law(pre, post, e): return all((not pre(s0)) or post(s0, e(s0)) for s0 in START)\n",
+    "# make the brute force genuinely quantify over Q for a subset of cases, to validate the \"worst Q = allowed\" shortcut\n",
+    "def refines_all_Q(pre, post, e):\n",
+    "    for s0 in START:\n",
+    "        for mask in range(1 << len(BOXX)):\n",
+    "            Q = {BOXX[k] for k in range(len(BOXX)) if mask >> k & 1}\n",
+    "            wp_spec = pre(s0) and all((xp in Q) for xp in BOXX if post(s0, xp))\n",
+    "            wp_impl = e(s0) in Q\n",
+    "            if wp_spec and not wp_impl: return False\n",
+    "    return True\n",
+    "random.seed(5)\n",
+    "tests = []\n",
+    "for _ in range(60):\n",
+    "    a, b, c = random.randint(-2, 2), random.randint(-2, 2), random.randint(1, 3)\n",
+    "    pre = (lambda a: lambda s0: s0 >= a)(a)\n",
+    "    post = (lambda b, c: lambda s0, xp: abs(xp - (s0 + b)) <= c - 1)(b, c)         # \"x ends within c-1 of x0+b\"\n",
+    "    e = (lambda b: lambda s0: s0 + b)(b)\n",
+    "    tests.append((pre, post, e))\n",
+    "# also deliberately failing implementations\n",
+    "for _ in range(60):\n",
+    "    a, b, c = random.randint(-2, 2), random.randint(-2, 2), random.randint(1, 2)\n",
+    "    pre = (lambda a: lambda s0: s0 >= a)(a)\n",
+    "    post = (lambda b, c: lambda s0, xp: abs(xp - (s0 + b)) <= c - 1)(b, c)\n",
+    "    k = random.randint(-1, 1)\n",
+    "    e = (lambda b, k: lambda s0: s0 + b + k)(b, k)\n",
+    "    tests.append((pre, post, e))\n",
+    "agree = 0; n_true = 0\n",
+    "for pre, post, e in tests[:40] + tests[60:100]:\n",
+    "    r1 = refines_all_Q(pre, post, e); r2 = refines_bruteforce(pre, post, e); r3 = law(pre, post, e)\n",
+    "    assert r1 == r2 == r3\n",
+    "    agree += 1; n_true += r1\n",
+    "print(f'{agree} random (spec, assignment) pairs: all-Q brute force == worst-Q shortcut == law pre => post[e/x]; refined in {n_true}, not refined in {agree-n_true}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cc1bc3fa",
+   "metadata": {},
+   "source": [
+    "## Summary of the checks\n",
+    "* `wp` by substitution == semantic wp on 14,700 (program, state) pairs (0 mismatches).\n",
+    "* Valid VCs (with variant) => triple holds and loop terminates (40 box states, 0 violations); 4 failure modes classified.\n",
+    "* Lean: `search_correct` proved (axioms `propext`, `Quot.sound`); the off-by-one variant is rejected and `#eval` produces a concrete violation.\n",
+    "* Overflow: VC over Z true, over `UInt8` false for exactly `#{lo+hi >= 256}` = 16,384 pairs, matching Python and Lean `#eval` counts.\n",
+    "* Refinement law == brute force over all 2^11 postconditions on 80 random pairs."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "bca6ef85",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:55:42.522453Z",
+     "iopub.status.busy": "2026-10-14T09:55:42.522294Z",
+     "iopub.status.idle": "2026-10-14T09:55:42.789460Z",
+     "shell.execute_reply": "2026-10-14T09:55:42.788381Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "L.close()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

code/notebooks/0304-failed-proof-to-bug.ipynb

@@ -0,0 +1,1006 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "4f055b7b",
+   "metadata": {},
+   "source": [
+    "# 0304 - From failed proof to bug: counterexample search and the soundness / completeness asymmetry\n",
+    "\n",
+    "Scratch-pad for note `0304 - Failed Proof to Bug.pdf`. Naive first, efficient second, assert agreement.\n",
+    "\n",
+    "* **A. Random testing with shrinking** (a 60-line QuickCheck). NAIVE = exhaustive enumeration of the bounded domain.\n",
+    "  EFFICIENT = random sampling + shrinking. Assert: the exact failure probability `p` (by enumeration) predicts the\n",
+    "  Monte-Carlo miss probability `(1-p)^N` and the mean time-to-failure `1/p`; shrinking reaches the global minimum.\n",
+    "* **B. Small-scope hypothesis**: when does enumeration beat random sampling? Closed forms, checked by Monte Carlo.\n",
+    "* **C. Lean**: `decide`, `slim_check` (this Mathlib pin's name for Plausible), `omega` failure messages; a counterexample turned into a\n",
+    "  *kernel-checked refutation*; and the cases where a failed proof / counterexample is **not** a bug.\n",
+    "\n",
+    "No GPU. Lean: ONE worker through `align.verify.lean`, nice 15, after the MemAvailable>=45 GB check."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "9540a401",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:29.990834Z",
+     "iopub.status.busy": "2026-10-14T09:54:29.990642Z",
+     "iopub.status.idle": "2026-10-14T09:54:29.997582Z",
+     "shell.execute_reply": "2026-10-14T09:54:29.996880Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
+    "import itertools, random, math, sys, collections\n",
+    "random.seed(2026)\n",
+    "\n",
+    "# ---- the code under test: insertion sort with a duplicate-handling bug -------------------------------------------\n",
+    "def isort_bug(xs):\n",
+    "    out = []\n",
+    "    for v in xs:\n",
+    "        k = 0\n",
+    "        while k < len(out) and out[k] < v: k += 1\n",
+    "        if k < len(out) and out[k] == v: continue          # BUG: drops the duplicate\n",
+    "        out.insert(k, v)\n",
+    "    return out\n",
+    "def isort_ok(xs):\n",
+    "    out = []\n",
+    "    for v in xs:\n",
+    "        k = 0\n",
+    "        while k < len(out) and out[k] < v: k += 1\n",
+    "        out.insert(k, v)\n",
+    "    return out\n",
+    "def prop_sort(f):\n",
+    "    def prop(xs):\n",
+    "        ys = f(xs)\n",
+    "        return all(a <= b for a, b in zip(ys, ys[1:])) and collections.Counter(ys) == collections.Counter(xs)\n",
+    "    return prop\n",
+    "\n",
+    "# Domain: lists of length 0..6 over {-2..2}; generator: length uniform in 0..6, entries uniform.\n",
+    "V = list(range(-2, 3)); MAXLEN = 6\n",
+    "def gen(rng):\n",
+    "    return [rng.choice(V) for _ in range(rng.randint(0, MAXLEN))]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "24aac6bf",
+   "metadata": {},
+   "source": [
+    "## A1. NAIVE: exhaustive enumeration gives the exact failure probability under the generator\n",
+    "`P(list) = (1/7) * 5^-len`. Sum the probability of the failing lists: this `p` is the per-test detection probability."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "f8691e4c",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:29.999823Z",
+     "iopub.status.busy": "2026-10-14T09:54:29.999658Z",
+     "iopub.status.idle": "2026-10-14T09:54:30.269920Z",
+     "shell.execute_reply": "2026-10-14T09:54:30.269084Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "domain size 19531; failing inputs 19205; exact per-test failure probability p = 0.498514\n"
+     ]
+    }
+   ],
+   "source": [
+    "prop_bug, prop_ok = prop_sort(isort_bug), prop_sort(isort_ok)\n",
+    "fail_prob = 0.0; fails = []\n",
+    "for L_ in range(MAXLEN + 1):\n",
+    "    for xs in itertools.product(V, repeat=L_):\n",
+    "        xs = list(xs)\n",
+    "        assert prop_ok(xs)                                   # the reference implementation passes everywhere\n",
+    "        if not prop_bug(xs):\n",
+    "            fails.append(xs); fail_prob += (1 / (MAXLEN + 1)) * 5.0 ** (-L_)\n",
+    "n_total = sum(5 ** L_ for L_ in range(MAXLEN + 1))\n",
+    "print(f'domain size {n_total}; failing inputs {len(fails)}; exact per-test failure probability p = {fail_prob:.6f}')\n",
+    "p = fail_prob"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ee87ac91",
+   "metadata": {},
+   "source": [
+    "## A2. EFFICIENT: random testing. Derivation (note 0304 section 3): `P(no failure in N independent tests) = (1-p)^N`;\n",
+    "time to first failure is Geometric(p) with mean `1/p`. Check both by Monte Carlo (20,000 runs)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "c36a5bed",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:30.271520Z",
+     "iopub.status.busy": "2026-10-14T09:54:30.271369Z",
+     "iopub.status.idle": "2026-10-14T09:54:30.557635Z",
+     "shell.execute_reply": "2026-10-14T09:54:30.556905Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "mean tests to first failure: empirical 1.992 +- 0.010 ; theory 1/p = 2.006 ; |diff| = 0.014 (1.43 se)\n",
+      "   N  P(miss) empirical    (1-p)^N      z\n",
+      "   1             0.5030     0.5015   0.41\n",
+      "   5             0.0297     0.0317   1.59\n",
+      "  10             0.0008     0.0010   0.70\n",
+      "  20             0.0000     0.0000   0.14\n",
+      "  40             0.0000     0.0000   0.00\n",
+      "  80             0.0000     0.0000   0.00\n",
+      "worst |z| = 1.59 (< 4: the geometric law (1-p)^N holds)\n",
+      "  N for miss-probability <= 0.1: 4\n",
+      "  N for miss-probability <= 0.01: 7\n",
+      "  N for miss-probability <= 0.001: 11\n"
+     ]
+    }
+   ],
+   "source": [
+    "rng = random.Random(7)\n",
+    "RUNS = 20000\n",
+    "first_fail = []\n",
+    "for _ in range(RUNS):\n",
+    "    t = 0\n",
+    "    while True:\n",
+    "        t += 1\n",
+    "        if not prop_bug(gen(rng)): break\n",
+    "    first_fail.append(t)\n",
+    "mean_t = sum(first_fail) / RUNS\n",
+    "se_mean = (sum((t - mean_t) ** 2 for t in first_fail) / (RUNS - 1)) ** 0.5 / RUNS ** 0.5\n",
+    "print(f'mean tests to first failure: empirical {mean_t:.3f} +- {se_mean:.3f} ; theory 1/p = {1/p:.3f} ; |diff| = {abs(mean_t - 1/p):.3f} ({abs(mean_t-1/p)/se_mean:.2f} se)')\n",
+    "assert abs(mean_t - 1 / p) < 4 * se_mean\n",
+    "worst_z = 0.0\n",
+    "print(f'{\"N\":>4s} {\"P(miss) empirical\":>18s} {\"(1-p)^N\":>10s} {\"z\":>6s}')\n",
+    "for N in (1, 5, 10, 20, 40, 80):\n",
+    "    emp = sum(t > N for t in first_fail) / RUNS\n",
+    "    th = (1 - p) ** N\n",
+    "    se = max((th * (1 - th) / RUNS) ** 0.5, 1e-9)\n",
+    "    z = abs(emp - th) / se; worst_z = max(worst_z, z)\n",
+    "    print(f'{N:4d} {emp:18.4f} {th:10.4f} {z:6.2f}')\n",
+    "assert worst_z < 4\n",
+    "print('worst |z| =', round(worst_z, 2), '(< 4: the geometric law (1-p)^N holds)')\n",
+    "# tests needed for confidence 1 - delta:  N >= ln(delta) / ln(1-p)\n",
+    "for delta in (0.1, 0.01, 1e-3):\n",
+    "    print(f'  N for miss-probability <= {delta}: {math.ceil(math.log(delta) / math.log(1 - p))}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "06805f67",
+   "metadata": {},
+   "source": [
+    "### A2b. A rarer bug: the same law at p = 1/1024 (a \"needle\": the code is wrong on exactly one input out of K = 1024)\n",
+    "Uniform sampling over the K inputs; per-test detection probability is exactly `p = 1/K`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "a033314f",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:30.559759Z",
+     "iopub.status.busy": "2026-10-14T09:54:30.559600Z",
+     "iopub.status.idle": "2026-10-14T09:54:32.010396Z",
+     "shell.execute_reply": "2026-10-14T09:54:32.009675Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "needle p=1/1024: mean tests to failure 1015.6 +- 18.3 ; theory 1024\n",
+      "  N=  256  P(miss) empirical 0.7790  (1-p)^N 0.7787  z=0.04\n",
+      "  N= 1024  P(miss) empirical 0.3643  (1-p)^N 0.3677  z=0.38\n",
+      "  N= 2048  P(miss) empirical 0.1297  (1-p)^N 0.1352  z=0.89\n",
+      "  N= 3072  P(miss) empirical 0.0497  (1-p)^N 0.0497  z=0.01\n",
+      "  N= 4096  P(miss) empirical 0.0173  (1-p)^N 0.0183  z=0.39\n",
+      "=> to be 99% sure of seeing a 1-in-1024 bug you need 4714 random tests; testing is a *statistical* instrument.\n"
+     ]
+    }
+   ],
+   "source": [
+    "K2 = 1024; p2 = 1 / K2; RUNS2 = 3000\n",
+    "rng = random.Random(23)\n",
+    "needle = 777\n",
+    "ff = []\n",
+    "for _ in range(RUNS2):\n",
+    "    t = 0\n",
+    "    while True:\n",
+    "        t += 1\n",
+    "        if rng.randrange(K2) == needle: break\n",
+    "    ff.append(t)\n",
+    "m2 = sum(ff) / RUNS2; se2 = (sum((t - m2) ** 2 for t in ff) / (RUNS2 - 1)) ** 0.5 / RUNS2 ** 0.5\n",
+    "print(f'needle p=1/{K2}: mean tests to failure {m2:.1f} +- {se2:.1f} ; theory {K2}')\n",
+    "assert abs(m2 - K2) < 4 * se2\n",
+    "zmax = 0.0\n",
+    "for N in (256, 1024, 2048, 3072, 4096):\n",
+    "    emp = sum(t > N for t in ff) / RUNS2; th = (1 - p2) ** N\n",
+    "    z = abs(emp - th) / max((th * (1 - th) / RUNS2) ** 0.5, 1e-9); zmax = max(zmax, z)\n",
+    "    print(f'  N={N:5d}  P(miss) empirical {emp:.4f}  (1-p)^N {th:.4f}  z={z:.2f}')\n",
+    "assert zmax < 4\n",
+    "print('=> to be 99% sure of seeing a 1-in-1024 bug you need', math.ceil(math.log(0.01) / math.log(1 - p2)), 'random tests; testing is a *statistical* instrument.')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c1f9f5f8",
+   "metadata": {},
+   "source": [
+    "## A3. Shrinking: from a long counterexample to a minimal one\n",
+    "Candidates: delete one element; move one element one step towards 0. Accept the first that still fails; stop when none does.\n",
+    "Termination: each accepted step strictly decreases the measure `(len, sum |v|)` in the lexicographic (well-founded) order.\n",
+    "A local minimum is not automatically a global one; compare with the exhaustive global minimum."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "2c7b4fa0",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:32.011875Z",
+     "iopub.status.busy": "2026-10-14T09:54:32.011718Z",
+     "iopub.status.idle": "2026-10-14T09:54:32.226368Z",
+     "shell.execute_reply": "2026-10-14T09:54:32.225418Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "global minimal counterexamples (exhaustive): [[0, 0]]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2000 shrinks: always a failing, locally minimal list; GLOBAL minimum reached in 401/2000 = 0.201; mean length before shrinking 4.57\n",
+      "final shrunk values: {(-2, -2): 420, (2, 2): 408, (0, 0): 401, (1, 1): 388, (-1, -1): 383}\n",
+      "local minima that are NOT global (e.g. [1,1]: deleting gives [1] (passes), moving towards 0 gives [0,1] (passes)): [[1, 1], [-1, -1], [2, 2], [-2, -2]]\n"
+     ]
+    }
+   ],
+   "source": [
+    "def shrink_candidates(xs):\n",
+    "    for k in range(len(xs)): yield xs[:k] + xs[k + 1:]\n",
+    "    for k, v in enumerate(xs):\n",
+    "        if v != 0: yield xs[:k] + [v - (1 if v > 0 else -1)] + xs[k + 1:]\n",
+    "def measure(xs): return (len(xs), sum(abs(v) for v in xs))\n",
+    "def shrink(xs, prop):\n",
+    "    steps = 0\n",
+    "    while True:\n",
+    "        for c in shrink_candidates(xs):\n",
+    "            if not prop(c):\n",
+    "                assert measure(c) < measure(xs)               # strict decrease\n",
+    "                xs, steps = c, steps + 1; break\n",
+    "        else:\n",
+    "            return xs, steps\n",
+    "global_min = min(map(measure, fails)); gmin_lists = sorted([f for f in fails if measure(f) == global_min])\n",
+    "print('global minimal counterexamples (exhaustive):', gmin_lists)\n",
+    "hit = tot = 0; lens_before = []; finals = collections.Counter()\n",
+    "rng = random.Random(11)\n",
+    "for _ in range(2000):\n",
+    "    while True:\n",
+    "        xs = gen(rng)\n",
+    "        if not prop_bug(xs): break\n",
+    "    ys, _ = shrink(xs, prop_bug)\n",
+    "    assert not prop_bug(ys) and all(prop_bug(c) for c in shrink_candidates(ys))        # still failing AND locally minimal\n",
+    "    assert measure(ys) >= global_min\n",
+    "    lens_before.append(len(xs)); tot += 1; hit += measure(ys) == global_min; finals[tuple(ys)] += 1\n",
+    "print(f'2000 shrinks: always a failing, locally minimal list; GLOBAL minimum reached in {hit}/{tot} = {hit/tot:.3f}; mean length before shrinking {sum(lens_before)/tot:.2f}')\n",
+    "print('final shrunk values:', dict(sorted(finals.items(), key=lambda kv: -kv[1])))\n",
+    "print('local minima that are NOT global (e.g. [1,1]: deleting gives [1] (passes), moving towards 0 gives [0,1] (passes)):',\n",
+    "      [list(k) for k in finals if measure(list(k)) != global_min])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b80dcca8",
+   "metadata": {},
+   "source": [
+    "## B. Small-scope hypothesis: enumeration vs uniform random sampling\n",
+    "A \"needle\" bug at one input `x0` of a `K = 2^16`-element domain. Uniform random with `N` tests detects with probability `1-(1-1/K)^N`\n",
+    "if the needle is uniformly placed; enumeration `0,1,...,N-1` detects iff `x0 < N`. If the bug location is `x0 ~ Geometric(q)`\n",
+    "(bugs cluster at small inputs, the empirical small-scope hypothesis), enumeration detects with probability `1-(1-q)^N` while random sampling gets\n",
+    "`E[1-(1-1/K)^N]` = tiny. Both closed forms, checked by Monte Carlo."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "d538673e",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:32.227875Z",
+     "iopub.status.busy": "2026-10-14T09:54:32.227713Z",
+     "iopub.status.idle": "2026-10-14T09:54:32.772482Z",
+     "shell.execute_reply": "2026-10-14T09:54:32.771731Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "uniform needle:   enumeration 0.0036 (theory 0.0031)   random 0.0027 (theory 0.0030)\n",
+      "clustered needle: enumeration 0.9823 (theory 0.9824)   random theory 0.0030   -> enumeration wins by x322\n",
+      "32-bit needle, delta = 1e-3: N >= 2.967e+10 uniform tests (no test budget finds it; a proof or a solver does)\n"
+     ]
+    }
+   ],
+   "source": [
+    "K = 2 ** 16; N = 200; MC = 40000\n",
+    "rng = random.Random(5)\n",
+    "# uniform placement\n",
+    "enum_u = sum(rng.randrange(K) < N for _ in range(MC)) / MC\n",
+    "rand_u = sum(any(rng.randrange(K) == x0 for _ in range(N)) for x0 in (rng.randrange(K) for _ in range(4000))) / 4000\n",
+    "th_enum_u, th_rand_u = N / K, 1 - (1 - 1 / K) ** N\n",
+    "print(f'uniform needle:   enumeration {enum_u:.4f} (theory {th_enum_u:.4f})   random {rand_u:.4f} (theory {th_rand_u:.4f})')\n",
+    "# geometric placement with q = 0.02\n",
+    "q = 0.02\n",
+    "def geom(r):\n",
+    "    x = 0\n",
+    "    while r.random() > q: x += 1\n",
+    "    return min(x, K - 1)\n",
+    "enum_g = sum(geom(rng) < N for _ in range(MC)) / MC\n",
+    "th_enum_g = 1 - (1 - q) ** N\n",
+    "th_rand_g = 1 - (1 - 1 / K) ** N        # uniform sampling detects a needle anywhere with the same probability\n",
+    "print(f'clustered needle: enumeration {enum_g:.4f} (theory {th_enum_g:.4f})   random theory {th_rand_g:.4f}   -> enumeration wins by x{th_enum_g/th_rand_g:.0f}')\n",
+    "se = lambda p_: (p_ * (1 - p_) / MC) ** 0.5\n",
+    "assert abs(enum_g - th_enum_g) < 4 * se(th_enum_g) and abs(enum_u - th_enum_u) < 4 * se(th_enum_u) + 1e-3\n",
+    "assert abs(rand_u - th_rand_u) < 4 * (th_rand_u * (1 - th_rand_u) / 4000) ** 0.5 + 1e-3\n",
+    "# a 32-bit needle: uniform random needs ~ ln(1/delta) * 2^32 tests\n",
+    "print('32-bit needle, delta = 1e-3: N >=', f'{math.log(1e-3) / math.log(1 - 2.0**-32):.3e}', 'uniform tests (no test budget finds it; a proof or a solver does)')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ae0e1586",
+   "metadata": {},
+   "source": [
+    "## C. Lean: three tools, three kinds of answer\n",
+    "* `decide`: exhaustive over a *bounded decidable* statement => a **proof** or a definite refutation (complete within the bound, kernel-checked).\n",
+    "* `slim_check` (Plausible's ancestor on this pin): random testing => a **counterexample** or *nothing* (and when nothing, it closes the goal with `sorry`!).\n",
+    "* `omega`: a decision procedure for linear integer arithmetic => proof, or a failure message listing \"a possible counterexample\" **constraints**."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "f695dda3",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:32.774131Z",
+     "iopub.status.busy": "2026-10-14T09:54:32.773969Z",
+     "iopub.status.idle": "2026-10-14T09:54:37.428445Z",
+     "shell.execute_reply": "2026-10-14T09:54:37.427528Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[ok] no messages\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "{'env': 2}"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "sys.path.insert(0, os.path.abspath('.'))\n",
+    "from _lean_helper import Lean\n",
+    "import re\n",
+    "L = Lean()\n",
+    "L.show(r'''\n",
+    "def search (P : Nat → Bool) (len : Nat) (n : Nat) : Nat :=\n",
+    "  if n < len then (if P n then n else search P len (n + 1)) else n\n",
+    "termination_by len - n\n",
+    "def Post (P : Nat → Bool) (len r : Nat) : Prop :=\n",
+    "  r ≤ len ∧ (r = len ∨ P r = true) ∧ ∀ i, i < r → P i = false\n",
+    "instance (P : Nat → Bool) (len r : Nat) : Decidable (Post P len r) := by unfold Post; infer_instance\n",
+    "def searchBadF (P : Nat → Bool) (len : Nat) : Nat → Nat → Nat\n",
+    "  | 0, n => n\n",
+    "  | fuel+1, n => if n + 1 < len then (if P n then n else searchBadF P len fuel (n+1)) else n\n",
+    "''', keep=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ce429a6a",
+   "metadata": {},
+   "source": [
+    "### C1. A counterexample is a proof of the negation (certified bug)\n",
+    "1. `slim_check` finds a violating input for the buggy loop (random).\n",
+    "2. Instantiate it and let `decide` evaluate: the kernel certifies `not Post ...` for that instance.\n",
+    "3. Conclude `not (forall P len, Post ...)`: a kernel-checked refutation with **no axioms**."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "4e9cab8a",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:37.430013Z",
+     "iopub.status.busy": "2026-10-14T09:54:37.429851Z",
+     "iopub.status.idle": "2026-10-14T09:54:37.738292Z",
+     "shell.execute_reply": "2026-10-14T09:54:37.737351Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "--- 1. random search ---\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[error] line 1: ===================\n",
+      "Found problems!\n",
+      "k := 2\n",
+      "len := 2\n",
+      "issue: ⋯ does not hold\n",
+      "(0 shrinks)\n",
+      "-------------------\n",
+      "--- 2. kernel-certified instance, 3. refutation of the universal spec ---\n",
+      "[info] line 7: 'searchBadF_not_correct' does not depend on any axioms\n"
+     ]
+    }
+   ],
+   "source": [
+    "print('--- 1. random search ---')\n",
+    "r1 = L.show(\"example : ∀ (k len : ℕ), Post (fun i => i == k) len (searchBadF (fun i => i == k) len len 0) := by slim_check\")\n",
+    "assert any('Found problems' in m['data'] for m in r1['messages'])\n",
+    "print('--- 2. kernel-certified instance, 3. refutation of the universal spec ---')\n",
+    "r2 = L.show(r'''\n",
+    "theorem bug_witness : ¬ Post (fun _ => false) 10 (searchBadF (fun _ => false) 10 20 0) := by decide\n",
+    "theorem searchBadF_not_correct : ¬ (∀ (P : Nat → Bool) (len : Nat), Post P len (searchBadF P len len 0)) := fun h => by\n",
+    "  have := h (fun _ => false) 10\n",
+    "  revert this\n",
+    "  decide\n",
+    "#print axioms searchBadF_not_correct\n",
+    "''', keep=True)\n",
+    "assert any('does not depend on any axioms' in m['data'] for m in r2['messages'])\n",
+    "assert not [m for m in r2['messages'] if m['severity'] == 'error']"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "422be44f",
+   "metadata": {},
+   "source": [
+    "### C2. No counterexample is NOT a proof, and `slim_check` says so with `sorry`\n",
+    "On a TRUE statement `slim_check` reports success and **closes the goal with `sorry`**: `#print axioms` shows `sorryAx`, so the strict\n",
+    "verifier's axiom gate rejects it. Property testing can only produce counterexamples, never certificates."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "f4ff841e",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:37.740312Z",
+     "iopub.status.busy": "2026-10-14T09:54:37.740138Z",
+     "iopub.status.idle": "2026-10-14T09:54:38.143859Z",
+     "shell.execute_reply": "2026-10-14T09:54:38.143005Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[info] line 1: Success\n",
+      "[warning] line 1: declaration uses 'sorry'\n",
+      "[info] line 2: 'tt' depends on axioms: [sorryAx]\n",
+      "=> axioms of the \"slim_check proof\": 'tt' depends on axioms: [sorryAx]\n",
+      "--- the same statement is PROVED by omega (a decision procedure for linear arithmetic): ---\n",
+      "[info] line 2: 'tt2' depends on axioms: [propext, Quot.sound]\n"
+     ]
+    }
+   ],
+   "source": [
+    "r = L.show(\"theorem tt : ∀ a b : ℕ, b ≤ a → a - b + b = a := by slim_check\\n#print axioms tt\")\n",
+    "axs = [m['data'] for m in r['messages'] if 'depends on axioms' in m['data']]\n",
+    "assert axs and 'sorryAx' in axs[0]\n",
+    "print('=> axioms of the \"slim_check proof\":', axs[0])\n",
+    "print('--- the same statement is PROVED by omega (a decision procedure for linear arithmetic): ---')\n",
+    "r = L.show(\"theorem tt2 : ∀ a b : ℕ, b ≤ a → a - b + b = a := by omega\\n#print axioms tt2\")\n",
+    "assert not [m for m in r['messages'] if m['severity'] == 'error']"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c0e698ee",
+   "metadata": {},
+   "source": [
+    "### C3. A counterexample to the *statement* is not necessarily a bug in the *program*\n",
+    "`forall a b : N, a - b + b = a` is false (truncated subtraction): counterexample `a=0, b=2`. The program (`-`) is doing what `Nat.sub` does;\n",
+    "the **spec** is missing `b <= a` (note 0302, Lemma 1). Protocol step: check the counterexample against the *precondition*."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "797ea289",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:38.145619Z",
+     "iopub.status.busy": "2026-10-14T09:54:38.145413Z",
+     "iopub.status.idle": "2026-10-14T09:54:38.414542Z",
+     "shell.execute_reply": "2026-10-14T09:54:38.413549Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[error] line 1: ===================\n",
+      "Found problems!\n",
+      "a := 0\n",
+      "b := 1\n",
+      "issue: 1 = 0 does not hold\n",
+      "(0 shrinks)\n",
+      "-------------------\n",
+      "python model of N-subtraction agrees: 2 != 0\n"
+     ]
+    }
+   ],
+   "source": [
+    "r = L.show(\"example : ∀ a b : ℕ, a - b + b = a := by slim_check\")\n",
+    "assert any('Found problems' in m['data'] for m in r['messages'])\n",
+    "print('python model of N-subtraction agrees:', max(0, 0 - 2) + 2, '!= 0')\n",
+    "assert max(0, 0 - 2) + 2 == 2"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f63c9230",
+   "metadata": {},
+   "source": [
+    "### C4. A failed proof is not a bug: `omega` abstracts a nonlinear term and reports a *spurious* counterexample\n",
+    "`n <= n*n` is TRUE for all naturals. `omega` treats `n*n` as an atom `b` and answers with the constraints `a - b >= 1, a >= 0`; those\n",
+    "are satisfiable for an unconstrained atom, but `b = a*a` is not unconstrained."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "4540734b",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:38.416132Z",
+     "iopub.status.busy": "2026-10-14T09:54:38.415974Z",
+     "iopub.status.idle": "2026-10-14T09:54:38.447380Z",
+     "shell.execute_reply": "2026-10-14T09:54:38.446354Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[error] line 1: omega could not prove the goal:\n",
+      "a possible counterexample may satisfy the constraints\n",
+      "  a - b ≥ 1\n",
+      "  a ≥ 0\n",
+      "where\n",
+      " a := ↑n\n",
+      " b := ↑n * ↑n\n",
+      "exhaustive check that the \"possible counterexample\" does not exist for n < 10^5: True\n",
+      "[info] line 2: 'le_sq'' depends on axioms: [propext]\n"
+     ]
+    }
+   ],
+   "source": [
+    "r = L.show(\"example : ∀ n : ℕ, n ≤ n * n := by omega\")\n",
+    "assert any('omega could not prove the goal' in m['data'] for m in r['messages'])\n",
+    "print('exhaustive check that the \"possible counterexample\" does not exist for n < 10^5:', all(n <= n * n for n in range(10 ** 5)))\n",
+    "r = L.show(\"theorem le_sq' : ∀ n : ℕ, n ≤ n * n := fun n => Nat.le_mul_self n\\n#print axioms le_sq'\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5d6f5442",
+   "metadata": {},
+   "source": [
+    "### C4b. `omega` on this pin is not complete even for PURE linear integer arithmetic (no atoms)\n",
+    "The Lean docstring (`Init/Tactics.lean`, pin v4.9.0-rc1): *\"It is not yet a full decision procedure (no \"dark\" or \"grey\" shadows)\"*. Pugh's classic system\n",
+    "`27 <= 11x + 13y <= 45,  -10 <= 7x - 9y <= 4` has real solutions but NO integer solution. A complete procedure proves `False`; this `omega` answers\n",
+    "\"a possible counterexample may satisfy the constraints\" although none exists, and there are no atoms to blame."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "554118d7",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:38.449055Z",
+     "iopub.status.busy": "2026-10-14T09:54:38.448851Z",
+     "iopub.status.idle": "2026-10-14T09:54:38.475501Z",
+     "shell.execute_reply": "2026-10-14T09:54:38.474256Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[error] line 2: omega could not prove the goal:\n",
+      "a possible counterexample may satisfy the constraints\n",
+      "  27 ≤ 11*a + 13*b ≤ 45\n",
+      "  -10 ≤ 7*a - 9*b ≤ 4\n",
+      "where\n",
+      " a := x\n",
+      " b := y\n",
+      "real vertices: [('113/190', '299/190'), ('59/38', '29/38'), ('55/38', '85/38'), ('457/190', '271/190')]\n",
+      "integer points in the bounding box [-1,3]x[-1,3] (which contains the whole real polytope): []\n",
+      "=> the system has real solutions (the polytope is non-empty) but no integer solution, yet omega reports a \"possible counterexample\": a failure message is NOT a countermodel.\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "from fractions import Fraction as Fr\n",
+    "r = L.show(\"\"\"\n",
+    "example (x y : Int) (h1 : 27 ≤ 11*x+13*y) (h2 : 11*x+13*y ≤ 45) (h3 : -10 ≤ 7*x-9*y) (h4 : 7*x-9*y ≤ 4) : False := by omega\n",
+    "\"\"\")\n",
+    "assert any('possible counterexample' in m['data'] for m in r['messages'])\n",
+    "# the real polytope: vertices of  u = 11x+13y in {27,45},  w = 7x-9y in {-10,4}   (x = (9u+13w)/190, y = (7u-11w)/190)\n",
+    "verts = [(Fr(9*u + 13*w, 190), Fr(7*u - 11*w, 190)) for u in (27, 45) for w in (-10, 4)]\n",
+    "xs, ys = [v[0] for v in verts], [v[1] for v in verts]\n",
+    "lo_x, hi_x, lo_y, hi_y = int(min(xs)) - 1, int(max(xs)) + 1, int(min(ys)) - 1, int(max(ys)) + 1\n",
+    "pts = [(x, y) for x in range(lo_x, hi_x + 1) for y in range(lo_y, hi_y + 1) if 27 <= 11*x+13*y <= 45 and -10 <= 7*x-9*y <= 4]\n",
+    "print('real vertices:', [(str(a), str(b)) for a, b in verts])\n",
+    "print(f'integer points in the bounding box [{lo_x},{hi_x}]x[{lo_y},{hi_y}] (which contains the whole real polytope): {pts}')\n",
+    "assert pts == []\n",
+    "print('=> the system has real solutions (the polytope is non-empty) but no integer solution, yet omega reports a \"possible counterexample\": a failure message is NOT a countermodel.')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8a0fc84d",
+   "metadata": {},
+   "source": [
+    "### C5. A failed VC is not a bug: the counterexample state may be unreachable (continuation of 0303 A3)\n",
+    "Weaken the loop invariant to `n <= len`. The lemma `forall P len n, n <= len -> Post P len (search P len n)` is FALSE\n",
+    "(`slim_check` finds a counterexample, below). But the program starts at `n = 0`. We check **exhaustively in a scope** (all `P : Fin 4 -> Bool`,\n",
+    "`len <= 4`, `n <= len`): every VC-counterexample is a state the real loop never visits, and every state it does visit satisfies the lemma."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "fcb52aca",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:38.477338Z",
+     "iopub.status.busy": "2026-10-14T09:54:38.477168Z",
+     "iopub.status.idle": "2026-10-14T09:54:39.860566Z",
+     "shell.execute_reply": "2026-10-14T09:54:39.859669Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[error] line 1: ===================\n",
+      "Found problems!\n",
+      "P := [0 ↦ true, 1 ↦ true, _ ↦ true]\n",
+      "len := 6\n",
+      "n := 5\n",
+      "guard: 5 ≤ 6\n",
+      "issue: ⋯ does not hold\n",
+      "(0 shrinks)\n",
+      "-------------------\n",
+      "scope: VC counterexamples (P,len,n): 111; of these reachable from n=0: 0; reachable states total 129, violating the lemma: 0\n",
+      "Lean count of VC counterexamples in the same scope: ['111']\n",
+      "=> every counterexample to the weakened verification condition is unreachable: spurious, NOT a bug in the program.\n",
+      "--- and the program-level statement (start n = 0) survives random testing (with the sorry caveat of C2): ---\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[info] line 1: Success\n",
+      "[warning] line 1: declaration uses 'sorry'\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "{'messages': [{'severity': 'info',\n",
+       "   'pos': {'line': 1, 'column': 76},\n",
+       "   'endPos': {'line': 1, 'column': 86},\n",
+       "   'data': 'Success'},\n",
+       "  {'severity': 'warning',\n",
+       "   'pos': {'line': 1, 'column': 0},\n",
+       "   'endPos': {'line': 1, 'column': 7},\n",
+       "   'data': \"declaration uses 'sorry'\"}],\n",
+       " 'env': 13}"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "r = L.show(\"example : ∀ (P : Nat → Bool) (len n : Nat), n ≤ len → Post P len (search P len n) := by slim_check\")\n",
+    "assert any('Found problems' in m['data'] for m in r['messages'])\n",
+    "def trace(P, length):\n",
+    "    n_, seen = 0, []\n",
+    "    while n_ <= length:\n",
+    "        seen.append(n_)\n",
+    "        if n_ >= length or P[n_]: break\n",
+    "        n_ += 1\n",
+    "    return seen\n",
+    "def search5(P, ln, n0):\n",
+    "    m = n0\n",
+    "    while m < ln and not P[m]: m += 1\n",
+    "    return m\n",
+    "def post5(P, ln, r_):\n",
+    "    return r_ <= ln and (r_ == ln or P[r_]) and all((not P[i]) for i in range(r_))\n",
+    "vc_cex = reach = reach_bad = 0; vc_cex_reachable = 0\n",
+    "for length in range(5):\n",
+    "    for bits in itertools.product([False, True], repeat=4):\n",
+    "        tr_ = set(trace(bits, length))\n",
+    "        for n_ in range(length + 1):\n",
+    "            ok = post5(bits, length, search5(bits, length, n_))\n",
+    "            if not ok:\n",
+    "                vc_cex += 1; vc_cex_reachable += (n_ in tr_)\n",
+    "            if n_ in tr_:\n",
+    "                reach += 1; reach_bad += (not ok)\n",
+    "print(f'scope: VC counterexamples (P,len,n): {vc_cex}; of these reachable from n=0: {vc_cex_reachable}; reachable states total {reach}, violating the lemma: {reach_bad}')\n",
+    "assert vc_cex > 0 and vc_cex_reachable == 0 and reach_bad == 0\n",
+    "r = L.run('''\n",
+    "#eval (List.range 5).foldl (fun acc len => acc + (List.range 16).foldl (fun a m =>\n",
+    "    let P : Nat → Bool := fun i => (m / (2 ^ i)) % 2 == 1\n",
+    "    a + (List.range (len + 1)).foldl (fun b n => if decide (Post P len (search P len n)) then b else b + 1) 0) 0) 0\n",
+    "''')\n",
+    "lean_vc = [m['data'].strip() for m in r['messages'] if m['severity'] == 'info']\n",
+    "print('Lean count of VC counterexamples in the same scope:', lean_vc)\n",
+    "assert int(lean_vc[0]) == vc_cex\n",
+    "print('=> every counterexample to the weakened verification condition is unreachable: spurious, NOT a bug in the program.')\n",
+    "print('--- and the program-level statement (start n = 0) survives random testing (with the sorry caveat of C2): ---')\n",
+    "L.show(\"example : ∀ (P : Nat → Bool) (len : Nat), Post P len (search P len 0) := by slim_check\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "92feb95d",
+   "metadata": {},
+   "source": [
+    "### C6. `decide` as model finding in a bounded scope (complete up to the bound)\n",
+    "Exhaustive over all predicates `P` on `Fin 4` (16 of them) and all `len <= 4`: the correct loop passes, the buggy one fails\n",
+    "on a specific witness. Compare with the same enumeration in Python."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "570ac396",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:39.862363Z",
+     "iopub.status.busy": "2026-10-14T09:54:39.862136Z",
+     "iopub.status.idle": "2026-10-14T09:54:39.961142Z",
+     "shell.execute_reply": "2026-10-14T09:54:39.960201Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Python exhaustive: 80 (P,len) scopes; correct loop violates 0, buggy loop violates 15; first witness ((False, False, False, False), 1)\n",
+      "Lean exhaustive violation counts (buggy, correct): ['15', '0']\n"
+     ]
+    }
+   ],
+   "source": [
+    "def py_search(P, length, n=0):\n",
+    "    while n < length:\n",
+    "        if P[n]: return n\n",
+    "        n += 1\n",
+    "    return n\n",
+    "def py_bad(P, length, n=0):\n",
+    "    while n + 1 < length:\n",
+    "        if P[n]: return n\n",
+    "        n += 1\n",
+    "    return n\n",
+    "def py_post(P, length, r):\n",
+    "    return r <= length and (r == length or (r < 4 and P[r])) and all((not P[i]) for i in range(r))\n",
+    "tot_ok = bad_cnt = 0\n",
+    "bad_witness = None\n",
+    "for length in range(5):\n",
+    "    for bits in itertools.product([False, True], repeat=4):\n",
+    "        assert py_post(bits, length, py_search(bits, length))\n",
+    "        if not py_post(bits, length, py_bad(bits, length)):\n",
+    "            bad_cnt += 1; bad_witness = bad_witness or (bits, length)\n",
+    "        tot_ok += 1\n",
+    "print(f'Python exhaustive: {tot_ok} (P,len) scopes; correct loop violates 0, buggy loop violates {bad_cnt}; first witness {bad_witness}')\n",
+    "r = L.run('''\n",
+    "#eval (List.range 5).foldl (fun acc len => acc + (List.range 16).foldl (fun a m =>\n",
+    "    let P : Nat → Bool := fun i => (m / (2 ^ i)) % 2 == 1\n",
+    "    if decide (Post P len (searchBadF P len len 0)) then a else a + 1) 0) 0\n",
+    "#eval (List.range 5).foldl (fun acc len => acc + (List.range 16).foldl (fun a m =>\n",
+    "    let P : Nat → Bool := fun i => (m / (2 ^ i)) % 2 == 1\n",
+    "    if decide (Post P len (search P len 0)) then a else a + 1) 0) 0\n",
+    "''')\n",
+    "cnts = [m['data'].strip() for m in r['messages'] if m['severity'] == 'info']\n",
+    "print('Lean exhaustive violation counts (buggy, correct):', cnts)\n",
+    "assert int(cnts[0]) == bad_cnt and int(cnts[1]) == 0"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1547dd78",
+   "metadata": {},
+   "source": [
+    "### C7. A reachable VC violation: what it does and does not prove (Lemma 4 of the note)\n",
+    "Program of 0303: `i := 0; s := 0; while i < n do s := s + inc(i); i := i + 1` with `inc(i) = i` (correct) or `i + 1` (buggy), precondition `n >= 0`,\n",
+    "post `Q: 2s = n(n-1)`, invariant `I: 0<=i<=n and 2s = i(i-1)`. The loop-head states `(n,i,s)` visited from `n >= 0` are `Reach`. Classify every violated VC:\n",
+    "* **preserve VC** (`I and i<n  =>  I[next]`): a reachable violating state `x` has a REACHABLE successor `x'` violating `I`: the annotation (or the code) is wrong, but `Q` may still hold;\n",
+    "* **exit VC** (`I and not(i<n)  =>  Q`): a reachable violating state is the state in which the loop actually EXITS: `Q` is violated by a real run: a bug witness (code or spec);\n",
+    "* unreachable violation: no conclusion."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "85eabecd",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:39.962887Z",
+     "iopub.status.busy": "2026-10-14T09:54:39.962663Z",
+     "iopub.status.idle": "2026-10-14T09:54:39.984403Z",
+     "shell.execute_reply": "2026-10-14T09:54:39.983516Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "A weak I, right code, right Q                  preserve-viol    0 (reachable   0)  exit-viol  504 (reachable   0)  Q holds on every run: True\n",
+      "B right I, right code, WRONG Q (n(n+1))        preserve-viol    0 (reachable   0)  exit-viol    7 (reachable   7)  Q holds on every run: False\n",
+      "C right I, BUGGY code, right Q                 preserve-viol   28 (reachable   7)  exit-viol    0 (reachable   0)  Q holds on every run: False\n",
+      "D too-strong I (i<=5), right code, right Q     preserve-viol    2 (reachable   2)  exit-viol    0 (reachable   0)  Q holds on every run: True\n",
+      "=> reachable preserve-violation: I is not an invariant (D: but Q still holds); reachable exit-violation: Q fails in a real run (B).\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "def run_head_states(n, inc):\n",
+    "    st, seen = (n, 0, 0), []\n",
+    "    while True:\n",
+    "        seen.append(st); nn, ii, ss = st\n",
+    "        if not ii < nn: return seen\n",
+    "        st = (nn, ii + 1, ss + inc(ii))\n",
+    "Pn = range(0, 8)\n",
+    "Ibox = [(n_, i_, s_) for n_ in range(0, 8) for i_ in range(-1, 9) for s_ in range(-4, 60)]\n",
+    "Qf = lambda n_, i_, s_: 2 * s_ == n_ * (n_ - 1)\n",
+    "def analyse(label, inc, I, Q):\n",
+    "    reach = {st for n_ in Pn for st in run_head_states(n_, inc)}\n",
+    "    step = lambda st: (st[0], st[1] + 1, st[2] + inc(st[1]))\n",
+    "    pres = [st for st in Ibox if I(*st) and st[1] < st[0] and not I(*step(st))]\n",
+    "    exits = [st for st in Ibox if I(*st) and not st[1] < st[0] and not Q(*st)]\n",
+    "    pres_r = [st for st in pres if st in reach]; exits_r = [st for st in exits if st in reach]\n",
+    "    # (a) a reachable preserve-violation has a reachable successor that violates I\n",
+    "    assert all(step(st) in reach and not I(*step(st)) for st in pres_r)\n",
+    "    # (b) a reachable exit-violation is the final state of a real run, and it violates Q\n",
+    "    finals = {run_head_states(n_, inc)[-1] for n_ in Pn}\n",
+    "    assert all(st in finals and not Q(*st) for st in exits_r)\n",
+    "    post_ok = all(Q(*run_head_states(n_, inc)[-1]) for n_ in Pn)\n",
+    "    print(f'{label:46s} preserve-viol {len(pres):4d} (reachable {len(pres_r):3d})  exit-viol {len(exits):4d} (reachable {len(exits_r):3d})  Q holds on every run: {post_ok}')\n",
+    "    return len(pres), len(pres_r), len(exits), len(exits_r), post_ok\n",
+    "good = lambda i: i; buggy = lambda i: i + 1\n",
+    "I_ok  = lambda n_, i_, s_: 0 <= i_ <= n_ and 2 * s_ == i_ * (i_ - 1)\n",
+    "I_weak = lambda n_, i_, s_: 0 <= i_ <= n_\n",
+    "I_tooStrong = lambda n_, i_, s_: 0 <= i_ <= n_ and 2 * s_ == i_ * (i_ - 1) and i_ <= 5\n",
+    "Q_wrong = lambda n_, i_, s_: 2 * s_ == n_ * (n_ + 1)\n",
+    "A = analyse('A weak I, right code, right Q', good, I_weak, Qf)\n",
+    "B = analyse('B right I, right code, WRONG Q (n(n+1))', good, I_ok, Q_wrong)\n",
+    "C = analyse('C right I, BUGGY code, right Q', buggy, I_ok, Qf)\n",
+    "D = analyse('D too-strong I (i<=5), right code, right Q', good, I_tooStrong, Qf)\n",
+    "assert A[2] > 0 and A[3] == 0                      # exit violations exist, all unreachable: nothing follows (the 0304 C5 story)\n",
+    "assert B[3] > 0 and not B[4]                       # reachable exit violation <=> Q really fails: a bug witness (here a wrong spec)\n",
+    "assert C[1] > 0 and not C[4]                       # reachable preserve violation in the buggy program\n",
+    "assert D[1] > 0 and D[4]                           # reachable preserve violation, I is not an invariant, yet Q holds on every run\n",
+    "print('=> reachable preserve-violation: I is not an invariant (D: but Q still holds); reachable exit-violation: Q fails in a real run (B).')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2da08f1c",
+   "metadata": {},
+   "source": [
+    "## The protocol, executed on one example\n",
+    "1. (counterexample) `slim_check`/`decide` gives `P = const false, len = 10` (or the shrunk `k=2, len=1`).\n",
+    "2. (precondition) Is it admissible: `Post` has no precondition other than `len`/`P` arbitrary -> yes.\n",
+    "3. (reachable) the loop is started at `n=0`: the witness is a top-level input, not an internal state -> reachable.\n",
+    "4. (evaluate) `#eval searchBadF (fun _ => false) 10 20 0 = 9`, spec demands `10`.\n",
+    "5. (intent) does the spec say what we want? Dafny's LinearSearch: `n = len` when nothing matches -> yes.\n",
+    "6. (certify) `searchBadF_not_correct` is a kernel-checked theorem with no axioms."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "1eb263c7",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:39.986228Z",
+     "iopub.status.busy": "2026-10-14T09:54:39.986009Z",
+     "iopub.status.idle": "2026-10-14T09:54:40.253611Z",
+     "shell.execute_reply": "2026-10-14T09:54:40.252685Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "L.close()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

code/notebooks/0305-spec-risk.ipynb

@@ -0,0 +1,789 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "9bfd3f75",
+   "metadata": {},
+   "source": [
+    "# 0305 - Spec generation risk: the proof of the wrong theorem\n",
+    "\n",
+    "Scratch-pad for note `0305 - Spec Generation Risk.pdf`. Naive first, efficient second, assert agreement.\n",
+    "\n",
+    "* **A. Spec validation by tests, VERINA-style.** NAIVE = decide soundness/completeness of a candidate post-condition by EXHAUSTIVE\n",
+    "  comparison with the ground truth over a bounded domain. EFFICIENT = decide it from positive / negative tests only (the reduction\n",
+    "  of VERINA section 4.1). Assert: the test-based verdict is never wrong in the \"refuted\" direction, and equals the exhaustive one with full tests.\n",
+    "* **A2. Mutation testing of a spec**: kill probability `1-(1-kappa)^m` per mutant, checked by Monte Carlo.\n",
+    "* **B. Lean**: a \"verified\" function that does nothing useful; a spec that leaks the implementation; a vacuous statement; and the\n",
+    "  **statement-rewrite** finding of this repository (docs/P2.md item 18) reproduced against four gates, using `align.verify.lean` (<= 2 workers, nice 15).\n",
+    "* **C. Risk model**: `P(spec wrong | proof accepted)` by Bayes, checked by Monte Carlo.\n",
+    "\n",
+    "No GPU."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "6dc6c1b4",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:42.902895Z",
+     "iopub.status.busy": "2026-10-14T09:54:42.902750Z",
+     "iopub.status.idle": "2026-10-14T09:54:42.912430Z",
+     "shell.execute_reply": "2026-10-14T09:54:42.911655Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "40 lists; 1600 input/output pairs\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
+    "import itertools, random, math, sys, collections\n",
+    "random.seed(305)\n",
+    "\n",
+    "# ---- domain: lists of length <= 3 over {0,1,2} --------------------------------------------------------------\n",
+    "V = (0, 1, 2)\n",
+    "LISTS = [tuple(t) for n in range(4) for t in itertools.product(V, repeat=n)]\n",
+    "PAIRS = [(x, y) for x in LISTS for y in LISTS]\n",
+    "print(len(LISTS), 'lists;', len(PAIRS), 'input/output pairs')\n",
+    "\n",
+    "# atoms an LLM-written spec is built from\n",
+    "srt     = lambda y: all(a <= b for a, b in zip(y, y[1:]))\n",
+    "srt_str = lambda y: all(a <  b for a, b in zip(y, y[1:]))\n",
+    "srt_rev = lambda y: all(a >= b for a, b in zip(y, y[1:]))\n",
+    "perm    = lambda x, y: collections.Counter(x) == collections.Counter(y)\n",
+    "len_eq  = lambda x, y: len(x) == len(y)\n",
+    "sum_eq  = lambda x, y: sum(x) == sum(y)\n",
+    "subset  = lambda x, y: set(y) <= set(x)\n",
+    "leak    = lambda x, y: y == tuple(sorted(x))\n",
+    "\n",
+    "Q = lambda x, y: srt(y) and perm(x, y)                       # ground-truth post-condition (precondition: True)\n",
+    "CANDIDATES = {\n",
+    " 'Q0  sorted & perm (truth)':           lambda x, y: srt(y) and perm(x, y),\n",
+    " 'S1  sorted only':                     lambda x, y: srt(y),\n",
+    " 'S2  sorted & perm & len<=2':          lambda x, y: srt(y) and perm(x, y) and len(y) <= 2,\n",
+    " 'S3  sorted & same len & same sum':    lambda x, y: srt(y) and len_eq(x, y) and sum_eq(x, y),\n",
+    " 'S4  strictly sorted & perm':          lambda x, y: srt_str(y) and perm(x, y),\n",
+    " 'S5  sorted & y == sorted(x) (leaks)': lambda x, y: srt(y) and leak(x, y),\n",
+    " 'S6  perm only':                       lambda x, y: perm(x, y),\n",
+    " 'S7  sorted & same len & y subset x':  lambda x, y: srt(y) and len_eq(x, y) and subset(x, y),\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e791e799",
+   "metadata": {},
+   "source": [
+    "## A1. Soundness / completeness of a candidate spec\n",
+    "VERINA (arXiv 2505.23135 section 4.1): with pre-condition `P = True`, `Q^` is **sound** iff `forall x y. Q^(x,y) => Q(x,y)` (it accepts only correct outputs)\n",
+    "and **complete** iff `forall x y. Q(x,y) => Q^(x,y)` (it accepts every correct output).\n",
+    "NAIVE: check both over all 1600 pairs. EFFICIENT (their reduction): soundness needs only the *negative* tests (pairs with `not Q`), completeness only the *positive* ones."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "35f1c93e",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:42.914044Z",
+     "iopub.status.busy": "2026-10-14T09:54:42.913844Z",
+     "iopub.status.idle": "2026-10-14T09:54:43.013138Z",
+     "shell.execute_reply": "2026-10-14T09:54:43.012157Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "positive pairs: 40  negative pairs: 1560\n",
+      "candidate                              sound complete  accepted-negatives rejected-positives\n",
+      "Q0  sorted & perm (truth)              True  True                   0                 0\n",
+      "S1  sorted only                        False True                 760                 0\n",
+      "S2  sorted & perm & len<=2             True  False                  0                27\n",
+      "S3  sorted & same len & same sum       False True                  22                 0\n",
+      "S4  strictly sorted & perm             True  False                  0                24\n",
+      "S5  sorted & y == sorted(x) (leaks)    True  True                   0                 0\n",
+      "S6  perm only                          False True                  72                 0\n",
+      "S7  sorted & same len & y subset x     False True                 120                 0\n",
+      "definition == test-based verdict for all 8 candidates (exact).\n"
+     ]
+    }
+   ],
+   "source": [
+    "POS = [(x, y) for x, y in PAIRS if Q(x, y)]\n",
+    "NEG = [(x, y) for x, y in PAIRS if not Q(x, y)]\n",
+    "print('positive pairs:', len(POS), ' negative pairs:', len(NEG))\n",
+    "table = {}\n",
+    "for name, Qh in CANDIDATES.items():\n",
+    "    # NAIVE: the definition, over every pair\n",
+    "    sound_def    = all((not Qh(x, y)) or Q(x, y) for x, y in PAIRS)\n",
+    "    complete_def = all((not Q(x, y)) or Qh(x, y) for x, y in PAIRS)\n",
+    "    # EFFICIENT: tests only (full suites here) -- soundness: Q^ rejects every negative test; completeness: Q^ accepts every positive test\n",
+    "    a = sum(1 for x, y in NEG if Qh(x, y))                  # accepted negatives  = soundness violations\n",
+    "    b = sum(1 for x, y in POS if not Qh(x, y))              # rejected positives  = completeness violations\n",
+    "    sound_t, complete_t = (a == 0), (b == 0)\n",
+    "    assert (sound_def, complete_def) == (sound_t, complete_t)\n",
+    "    table[name] = (sound_t, complete_t, a, b)\n",
+    "print(f'{\"candidate\":38s} sound complete  accepted-negatives rejected-positives')\n",
+    "for name, (s_, c_, a_, b_) in table.items(): print(f'{name:38s} {str(s_):5s} {str(c_):8s} {a_:15d} {b_:17d}')\n",
+    "assert table['Q0  sorted & perm (truth)'][:2] == (True, True) and table['S5  sorted & y == sorted(x) (leaks)'][:2] == (True, True)\n",
+    "print('definition == test-based verdict for all', len(CANDIDATES), 'candidates (exact).')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8c5cdb33",
+   "metadata": {},
+   "source": [
+    "### A1b. A wrong *program* passes a wrong *spec*  (the proof of the wrong theorem, at toy scale)\n",
+    "Programs `f : list -> list`. \"f satisfies spec S\" means `forall x. S(x, f(x))`. Truth = f is a correct sort."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "2302af24",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:43.014728Z",
+     "iopub.status.busy": "2026-10-14T09:54:43.014577Z",
+     "iopub.status.idle": "2026-10-14T09:54:43.023261Z",
+     "shell.execute_reply": "2026-10-14T09:54:43.022309Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "program                                 Q0   S1   S2   S3   S4   S5   S6   S7\n",
+      "sorted (correct)                      PASS PASS   .  PASS   .  PASS PASS PASS\n",
+      "empty list                              .  PASS   .    .    .    .    .    . \n",
+      "identity                                .    .    .    .    .    .  PASS   . \n",
+      "sorted dedup                            .  PASS   .    .    .    .    .    . \n",
+      "sorted, replace each v by min(v,1)      .  PASS   .    .    .    .    .    . \n",
+      "sorted descending                       .    .    .    .    .    .  PASS   . \n",
+      "sort then zero the first element        .  PASS   .    .    .    .    .    . \n",
+      "wrong programs accepted by each spec: {'S1  sorted only': 4, 'S6  perm only': 2}\n"
+     ]
+    }
+   ],
+   "source": [
+    "PROGRAMS = {\n",
+    " 'sorted (correct)':                   lambda x: tuple(sorted(x)),\n",
+    " 'empty list':                         lambda x: (),\n",
+    " 'identity':                           lambda x: x,\n",
+    " 'sorted dedup':                       lambda x: tuple(sorted(set(x))),\n",
+    " 'sorted, replace each v by min(v,1)': lambda x: tuple(sorted(min(v, 1) for v in x)),\n",
+    " 'sorted descending':                  lambda x: tuple(sorted(x, reverse=True)),\n",
+    " 'sort then zero the first element':   lambda x: tuple(sorted(x))[:0] + ((0,) + tuple(sorted(x))[1:] if x else ()),\n",
+    "}\n",
+    "def passes(spec, f): return all(spec(x, f(x)) for x in LISTS)\n",
+    "print(f'{\"program\":38s}' + ' '.join(f'{n.split()[0]:>4s}' for n in CANDIDATES))\n",
+    "truth_pass = {pn: passes(Q, f) for pn, f in PROGRAMS.items()}\n",
+    "false_accepts = collections.Counter()\n",
+    "for pn, f in PROGRAMS.items():\n",
+    "    row = []\n",
+    "    for sn, S in CANDIDATES.items():\n",
+    "        ok = passes(S, f); row.append(ok)\n",
+    "        if ok and not truth_pass[pn]: false_accepts[sn] += 1\n",
+    "    print(f'{pn:38s}' + ' '.join(f'{\"PASS\" if o else \"  . \":>4s}' for o in row))\n",
+    "assert false_accepts['S1  sorted only'] >= 3            # empty, dedup, min-clamp, zero-first, ... all pass the weak spec\n",
+    "assert false_accepts['Q0  sorted & perm (truth)'] == 0\n",
+    "print('wrong programs accepted by each spec:', dict(false_accepts))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4371ed45",
+   "metadata": {},
+   "source": [
+    "## A2. Mutation testing of a spec\n",
+    "Mutants of `Q0 = sorted(y) and perm(x,y)`: drop an atom, weaken/strengthen an atom. A test suite `T` of `m` pairs *kills* a mutant `M` iff some test has `M(x,y) != Q0(x,y)`.\n",
+    "Let `kappa(M) = #{(x,y): M != Q0} / 1600`. Uniform tests: `P(kill) = 1 - (1-kappa)^m`. `kappa = 0` means an *equivalent* mutant (unkillable)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "aeda6210",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:43.024893Z",
+     "iopub.status.busy": "2026-10-14T09:54:43.024737Z",
+     "iopub.status.idle": "2026-10-14T09:54:51.227959Z",
+     "shell.execute_reply": "2026-10-14T09:54:51.226999Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "mutant                              kappa  P(kill) m=10 theory  empirical     z\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "drop sorted                        0.0450               0.3690     0.3715  0.75\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "drop perm                          0.4750               0.9984     0.9986  0.50\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "sorted: <= -> <                    0.0150               0.1403     0.1448  1.84\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "sorted: <= -> >=                   0.0375               0.3176     0.3192  0.49\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "perm -> same length                0.1837               0.8687     0.8689  0.08\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "perm -> same sum                   0.0600               0.4614     0.4608  0.15\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "perm -> subset                     0.2231               0.9199     0.9199  0.01\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "add redundant len(x)==len(y)       0.0000               0.0000     0.0000  0.00\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "perm -> y == sorted(x)             0.0000               0.0000     0.0000  0.00\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "add len(y)<=2                      0.0169               0.1565     0.1540  0.97\n",
+      "equivalent (unkillable) mutants: ['add redundant len(x)==len(y)', 'perm -> y == sorted(x)']\n",
+      "mutation score with the exhaustive suite: 8/8 non-equivalent killed; worst |z| of the kill-probability law = 1.84\n"
+     ]
+    }
+   ],
+   "source": [
+    "MUTANTS = {\n",
+    " 'drop sorted':                lambda x, y: perm(x, y),\n",
+    " 'drop perm':                  lambda x, y: srt(y),\n",
+    " 'sorted: <= -> <':            lambda x, y: srt_str(y) and perm(x, y),\n",
+    " 'sorted: <= -> >=':           lambda x, y: srt_rev(y) and perm(x, y),\n",
+    " 'perm -> same length':        lambda x, y: srt(y) and len_eq(x, y),\n",
+    " 'perm -> same sum':           lambda x, y: srt(y) and sum_eq(x, y),\n",
+    " 'perm -> subset':             lambda x, y: srt(y) and subset(x, y),\n",
+    " 'add redundant len(x)==len(y)': lambda x, y: srt(y) and perm(x, y) and len_eq(x, y),     # equivalent: perm implies equal length\n",
+    " 'perm -> y == sorted(x)':     lambda x, y: srt(y) and leak(x, y),                          # equivalent\n",
+    " 'add len(y)<=2':              lambda x, y: srt(y) and perm(x, y) and len(y) <= 2,\n",
+    "}\n",
+    "kappa = {n: sum(1 for x, y in PAIRS if M(x, y) != Q(x, y)) / len(PAIRS) for n, M in MUTANTS.items()}\n",
+    "rng = random.Random(9); M_SIZE = 10; SUITES = 20000\n",
+    "print(f'{\"mutant\":32s} {\"kappa\":>8s} {\"P(kill) m=10 theory\":>20s} {\"empirical\":>10s} {\"z\":>5s}')\n",
+    "worst = 0.0\n",
+    "for n, M in MUTANTS.items():\n",
+    "    kills = 0\n",
+    "    for _ in range(SUITES):\n",
+    "        suite = [PAIRS[rng.randrange(len(PAIRS))] for _ in range(M_SIZE)]\n",
+    "        kills += any(M(x, y) != Q(x, y) for x, y in suite)\n",
+    "    emp = kills / SUITES; th = 1 - (1 - kappa[n]) ** M_SIZE\n",
+    "    se = max((th * (1 - th) / SUITES) ** 0.5, 1e-9); z = abs(emp - th) / se if th not in (0.0,) else 0.0\n",
+    "    worst = max(worst, z)\n",
+    "    print(f'{n:32s} {kappa[n]:8.4f} {th:20.4f} {emp:10.4f} {z:5.2f}')\n",
+    "assert worst < 4.0\n",
+    "equiv = [n for n, k in kappa.items() if k == 0]\n",
+    "print('equivalent (unkillable) mutants:', equiv)\n",
+    "assert set(equiv) == {'add redundant len(x)==len(y)', 'perm -> y == sorted(x)'}\n",
+    "# exhaustive suite kills every non-equivalent mutant\n",
+    "assert all(any(M(x, y) != Q(x, y) for x, y in PAIRS) for n, M in MUTANTS.items() if n not in equiv)\n",
+    "print(f'mutation score with the exhaustive suite: {(len(MUTANTS) - len(equiv))}/{len(MUTANTS) - len(equiv)} non-equivalent killed; worst |z| of the kill-probability law = {worst:.2f}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "76064a25",
+   "metadata": {},
+   "source": [
+    "## B. Lean: proof of the wrong theorem, four ways\n",
+    "ONE worker via `align.verify.lean` for the experiments; a second (the strict verifier's) for B4. Total <= 2 workers, nice 15."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "cde1b268",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:51.229911Z",
+     "iopub.status.busy": "2026-10-14T09:54:51.229754Z",
+     "iopub.status.idle": "2026-10-14T09:54:55.294102Z",
+     "shell.execute_reply": "2026-10-14T09:54:55.293131Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "sys.path.insert(0, os.path.abspath('.'))\n",
+    "from _lean_helper import Lean\n",
+    "L = Lean()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "46718286",
+   "metadata": {},
+   "source": [
+    "### B1. A weak specification is satisfied by a useless function (and Lean happily certifies it)\n",
+    "Intended spec: `Sorted /\\ Perm`. A weak spec: `Sorted` only. `badSort := fun _ => []` is \"verified\" against the weak spec; the full spec refutes it."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "b665be7a",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:55.295817Z",
+     "iopub.status.busy": "2026-10-14T09:54:55.295655Z",
+     "iopub.status.idle": "2026-10-14T09:54:55.352936Z",
+     "shell.execute_reply": "2026-10-14T09:54:55.352189Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[warning] line 2: unused variable `xs`\n",
+      "note: this linter can be disabled with `set_option linter.unusedVariables false`\n",
+      "[info] line 4: 'badSort_weak_spec' does not depend on any axioms\n",
+      "[info] line 8: 'badSort_not_perm' depends on axioms: [propext]\n"
+     ]
+    }
+   ],
+   "source": [
+    "r = L.show(r'''\n",
+    "def badSort (xs : List ℕ) : List ℕ := []\n",
+    "theorem badSort_weak_spec (xs : List ℕ) : List.Sorted (· ≤ ·) (badSort xs) := List.sorted_nil\n",
+    "#print axioms badSort_weak_spec\n",
+    "theorem badSort_not_perm : ¬ (∀ xs : List ℕ, List.Perm xs (badSort xs)) := fun h => by\n",
+    "  have := (h [1]).length_eq\n",
+    "  simp [badSort] at this\n",
+    "#print axioms badSort_not_perm\n",
+    "''', keep=True)\n",
+    "assert not [m for m in r.get('messages', []) if m['severity'] == 'error']"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5f267817",
+   "metadata": {},
+   "source": [
+    "### B2. A specification that leaks the implementation (CLEVER, arXiv 2505.13938, Fig. 2)\n",
+    "Computable spec = the recursion itself: the proof is `rfl`-level, no mathematics. A relational (non-computable) spec forces an actual induction."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "f9513e46",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:55.354522Z",
+     "iopub.status.busy": "2026-10-14T09:54:55.354372Z",
+     "iopub.status.idle": "2026-10-14T09:54:55.422221Z",
+     "shell.execute_reply": "2026-10-14T09:54:55.421465Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[ok] no messages\n"
+     ]
+    }
+   ],
+   "source": [
+    "r = L.show(r'''\n",
+    "def myfib : ℕ → ℕ | 0 => 0 | 1 => 1 | (n+2) => myfib n + myfib (n+1)\n",
+    "-- (a) the spec IS the implementation\n",
+    "def spec_leak (impl : ℕ → ℕ) (n : ℕ) : Prop :=\n",
+    "  impl n = match n with | 0 => 0 | 1 => 1 | (n'+2) => impl n' + impl (n'+1)\n",
+    "theorem leak_proof : ∀ n, spec_leak myfib n := by\n",
+    "  intro n; match n with | 0 => rfl | 1 => rfl | n'+2 => rfl\n",
+    "-- (b) a relational spec: needs induction\n",
+    "inductive IsFib : ℕ → ℕ → Prop\n",
+    "  | zero : IsFib 0 0\n",
+    "  | one  : IsFib 1 1\n",
+    "  | step : ∀ n a b, IsFib n a → IsFib (n+1) b → IsFib (n+2) (a+b)\n",
+    "theorem myfib_isfib : ∀ n, IsFib n (myfib n) ∧ IsFib (n+1) (myfib (n+1)) := by\n",
+    "  intro n; induction n with\n",
+    "  | zero => exact ⟨IsFib.zero, IsFib.one⟩\n",
+    "  | succ k ih => exact ⟨ih.2, IsFib.step k _ _ ih.1 ih.2⟩\n",
+    "''', keep=True)\n",
+    "assert not [m for m in r.get('messages', []) if m['severity'] == 'error']"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ae047cf3",
+   "metadata": {},
+   "source": [
+    "### B3. Vacuous hypotheses, and a probe that detects them\n",
+    "`h0 : n < 5`, `h1 : n > 7` are contradictory; any conclusion follows. Probe: is `exists n < 100, hypotheses` satisfiable?  (A scoped *satisfiability witness* search; ProofGate-style negation probes are the dual.)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "a201ade3",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:55.423774Z",
+     "iopub.status.busy": "2026-10-14T09:54:55.423624Z",
+     "iopub.status.idle": "2026-10-14T09:54:55.506541Z",
+     "shell.execute_reply": "2026-10-14T09:54:55.505667Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[info] line 3: 'vac_stmt' depends on axioms: [propext, Quot.sound]\n",
+      "[info] line 4: false\n",
+      "[info] line 5: true\n",
+      "python probe agrees: the contradictory hypotheses have no witness in the scope; the consistent pair does.\n"
+     ]
+    }
+   ],
+   "source": [
+    "r = L.show(r'''\n",
+    "theorem vac_stmt (n : ℕ) (h₀ : n < 5) (h₁ : n > 7) : n = 12345 := by omega\n",
+    "#print axioms vac_stmt\n",
+    "#eval decide (∃ n < (100:ℕ), n < 5 ∧ n > 7)\n",
+    "#eval decide (∃ n < (100:ℕ), n < 10 ∧ n > 7)\n",
+    "''')\n",
+    "assert not [m for m in r.get('messages', []) if m['severity'] == 'error']\n",
+    "probe = lambda hyps: any(all(h(n_) for h in hyps) for n_ in range(100))\n",
+    "assert probe([lambda n_: n_ < 5, lambda n_: n_ > 7]) is False and probe([lambda n_: n_ < 10, lambda n_: n_ > 7]) is True\n",
+    "print('python probe agrees: the contradictory hypotheses have no witness in the scope; the consistent pair does.')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "94f56031",
+   "metadata": {},
+   "source": [
+    "### B4. The statement-rewrite finding (docs/P2.md section 8c item 18), reproduced against four gates\n",
+    "Reference: miniF2F `mathd_algebra_478` (`v = 65`). Candidates are complete files whose DISPLAYED theorem differs from the reference.\n",
+    "Gates: (1) **compile-only** (no error, no `sorry`); (2) **swap** (Goedel's `problem_check`: replace the displayed signature by the reference one and re-run the tactics);\n",
+    "(3) **strict** (token-equal signature + kernel statement bridge + axiom audit); (4) **defeq** (bridge only)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "dda78612",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:54:55.508277Z",
+     "iopub.status.busy": "2026-10-14T09:54:55.508117Z",
+     "iopub.status.idle": "2026-10-14T09:55:05.649327Z",
+     "shell.execute_reply": "2026-10-14T09:55:05.648446Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "theorem mathd_algebra_478 (b h v : ℝ) (h₀ : 0 < b ∧ 0 < h ∧ 0 < v) (h₁ : v = 1 / 3 * (b * h))\n",
+      "    (h₂ : b = 30) (h₃ : h = 13 / 2) : v = 65\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "candidate                    compile swap  strict defeq   strict-reasons / defeq-reasons\n",
+      "C0 honest                    True   True  True   True    [] / []\n",
+      "C1 displays v = 66           False  True  False  False   ['statement_mismatch'] / ['compile_error']\n",
+      "C2 weakened v ≤ 65           True   True  False  False   ['statement_mismatch'] / ['bridge_failed']\n",
+      "C3 extra false hypothesis    True   False False  False   ['statement_mismatch'] / ['bridge_failed']\n",
+      "C4 alpha-renamed binders     True   True  False  True    ['statement_mismatch'] / []\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "compile-only credits C2 (weaker) and C3 (vacuous) -- each is a different, easier theorem; swap credits C1 and C2 (the displayed statement is not the one proved) and refuses C3 only because its tactics name the extra hypothesis; strict rejects C1-C4; defeq accepts only the alpha-renaming.\n",
+      "accepted sets: {'compile': ['C0', 'C2', 'C3', 'C4'], 'swap': ['C0', 'C1', 'C2', 'C4'], 'strict': ['C0'], 'defeq': ['C0', 'C4']}\n",
+      "=> compile-only and swap are incomparable (C1 passes only swap, C3 only compile-only); defeq is stricter than both; strict is stricter than defeq: a partial order, not a chain.\n"
+     ]
+    }
+   ],
+   "source": [
+    "import json\n",
+    "from align.verify.lean import ReplPool, ReplLimits, StrictVerifier, StrictConfig, find_toolchain\n",
+    "rows = [json.loads(l) for l in open('/apps/development/rlhf/p2_runs/r1_baseline/minif2f.jsonl')]\n",
+    "ref = next(r_ for r_ in rows if r_['problem_id'] == 'mathd_algebra_478')['formal_statement']\n",
+    "head, tail = ref.split(':= by sorry')\n",
+    "sig = head[head.index('theorem'):]\n",
+    "print(sig.strip()); print()\n",
+    "PROOF = 'by rw [h₁, h₂, h₃]; norm_num'\n",
+    "SIGS = {\n",
+    " 'C0 honest':                 sig,\n",
+    " 'C1 displays v = 66':        sig.replace(': v = 65', ': v = 66'),\n",
+    " 'C2 weakened v ≤ 65':        sig.replace(': v = 65', ': v ≤ 65'),\n",
+    " 'C3 extra false hypothesis': sig.replace('(h₃ : h = 13 / 2)', '(h₃ : h = 13 / 2) (h₄ : v < 0)'),\n",
+    " 'C4 alpha-renamed binders':  sig.replace('(b h v : ℝ)', '(b₁ h v : ℝ)').replace('0 < b ∧', '0 < b₁ ∧').replace('(b * h)', '(b₁ * h)').replace('(h₂ : b =', '(h₂ : b₁ ='),\n",
+    "}\n",
+    "CODES = {k: head.replace(sig, s) + ':= ' + (PROOF if k != 'C3 extra false hypothesis' else 'by linarith [h₀.1, h₀.2.2, h₄]') + '\\n' for k, s in SIGS.items()}\n",
+    "\n",
+    "def gate_compile(code):\n",
+    "    code = '\\n'.join(l for l in code.split('\\n') if not l.startswith('import '))      # the worker already imported Mathlib\n",
+    "    r_ = L.run(code)\n",
+    "    errs = [m for m in r_.get('messages', []) if m['severity'] == 'error']\n",
+    "    sorry = r_.get('sorries') or any('sorry' in m['data'] for m in r_.get('messages', []))\n",
+    "    return (not errs) and (not sorry)\n",
+    "def gate_swap(code):\n",
+    "    proof = code.split(':=', 1)[1]                       # tactics after the displayed signature\n",
+    "    return gate_compile(head + ':=' + proof)             # re-run them against the ORIGINAL statement (Goedel's problem_check)\n",
+    "\n",
+    "pool = ReplPool.from_toolchain(find_toolchain(), n_workers=1, limits=ReplLimits(nice=15))\n",
+    "V_strict = StrictVerifier(pool, StrictConfig(statement_mode='strict'))\n",
+    "V_defeq = StrictVerifier(pool, StrictConfig(statement_mode='defeq'))\n",
+    "results = {}\n",
+    "for k, code in CODES.items():\n",
+    "    c = gate_compile(code); sw = gate_swap(code)\n",
+    "    rs = V_strict.verify(ref, code); rd = V_defeq.verify(ref, code)\n",
+    "    results[k] = (c, sw, rs.accepted, rd.accepted, rs.reasons, rd.reasons)\n",
+    "print(f'{\"candidate\":28s} compile swap  strict defeq   strict-reasons / defeq-reasons')\n",
+    "for k, (c, sw, s_, d_, rs_, rd_) in results.items():\n",
+    "    print(f'{k:28s} {str(c):7s}{str(sw):6s}{str(s_):7s}{str(d_):6s}  {list(rs_)} / {list(rd_)}')\n",
+    "pool.close()\n",
+    "# what the table must show (and is asserted):\n",
+    "assert results['C0 honest'][:4] == (True, True, True, True)\n",
+    "assert results['C1 displays v = 66'][0] is False and results['C1 displays v = 66'][1] is True     # swap harness credits a rewritten statement\n",
+    "assert results['C1 displays v = 66'][2] is False and results['C1 displays v = 66'][3] is False\n",
+    "assert results['C2 weakened v ≤ 65'][0] is True and results['C2 weakened v ≤ 65'][2] is False and results['C2 weakened v ≤ 65'][3] is False\n",
+    "assert results['C3 extra false hypothesis'][0] is True and results['C3 extra false hypothesis'][2] is False\n",
+    "assert results['C4 alpha-renamed binders'][2] is False and results['C4 alpha-renamed binders'][3] is True\n",
+    "print('\\ncompile-only credits C2 (weaker) and C3 (vacuous) -- each is a different, easier theorem; swap credits C1 and C2 (the displayed statement is not the one proved) and refuses C3 only because its tactics name the extra hypothesis; strict rejects C1-C4; defeq accepts only the alpha-renaming.')\n",
+    "\n",
+    "# Are the gates a CHAIN (each at least as strict as the previous)?  Compare the accepted sets over these five candidates.\n",
+    "acc = {g: {k for k, v in results.items() if v[idx]} for g, idx in (('compile', 0), ('swap', 1), ('strict', 2), ('defeq', 3))}\n",
+    "print('accepted sets:', {g: sorted(s_.split()[0] for s_ in v) for g, v in acc.items()})\n",
+    "assert not acc['compile'] <= acc['swap'] and not acc['swap'] <= acc['compile']      # compile-only vs swap: INCOMPARABLE (C1 only swap; C3 only compile)\n",
+    "assert acc['strict'] < acc['defeq'] and acc['defeq'] < acc['compile'] and acc['defeq'] < acc['swap']   # defeq is stricter than both, strict stricter than defeq\n",
+    "print('=> compile-only and swap are incomparable (C1 passes only swap, C3 only compile-only); defeq is stricter than both; strict is stricter than defeq: a partial order, not a chain.')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d02430be",
+   "metadata": {},
+   "source": [
+    "### B5. `autoImplicit`, the silent statement change (from 0302), as a spec-risk\n",
+    "See 0302: a typo becomes a fresh universally quantified variable. A one-line check at benchmark-build time: elaborate with `autoImplicit false`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "aaab2806",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:55:05.651602Z",
+     "iopub.status.busy": "2026-10-14T09:55:05.651417Z",
+     "iopub.status.idle": "2026-10-14T09:55:05.674529Z",
+     "shell.execute_reply": "2026-10-14T09:55:05.673620Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[warning] line 1: declaration uses 'sorry'\n",
+      "[info] line 2: @typo_stmt : ∀ {xx : ℕ} (x : ℕ), x + 1 > xx\n",
+      "[sorries] xx x : ℕ\n",
+      "⊢ x + 1 > xx\n",
+      "[error] line 2: unknown identifier 'xx'\n",
+      "[sorries] x : ℕ\n",
+      "⊢ x + 1 > sorryAx ℕ true\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "{'sorries': [{'proofState': 1,\n",
+       "   'pos': {'line': 2, 'column': 46},\n",
+       "   'goal': 'x : ℕ\\n⊢ x + 1 > sorryAx ℕ true',\n",
+       "   'endPos': {'line': 2, 'column': 51}}],\n",
+       " 'messages': [{'severity': 'error',\n",
+       "   'pos': {'line': 2, 'column': 37},\n",
+       "   'endPos': {'line': 2, 'column': 39},\n",
+       "   'data': \"unknown identifier 'xx'\"}],\n",
+       " 'env': 16}"
+      ]
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "L.show(\"theorem typo_stmt (x : ℕ) : x + 1 > xx := by sorry\\n#check @typo_stmt\")\n",
+    "L.show(\"set_option autoImplicit false in\\ntheorem typo_stmt2 (x : ℕ) : x + 1 > xx := by sorry\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ac95d158",
+   "metadata": {},
+   "source": [
+    "## C. Risk model: how likely is an accepted proof to be a proof of the wrong theorem?\n",
+    "Let a generated spec be wrong with probability `u`. A prover finds a proof with probability `a_c` for a correct spec and `a_w` for a wrong spec\n",
+    "(wrong = weaker/vacuous specs are easier: `r = a_w / a_c >= 1`). By Bayes\n",
+    "`P(wrong | accepted) = u a_w / (u a_w + (1-u) a_c) = u r / (u r + 1 - u)`. Check by Monte Carlo."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "7659f6ec",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:55:05.676725Z",
+     "iopub.status.busy": "2026-10-14T09:55:05.676540Z",
+     "iopub.status.idle": "2026-10-14T09:55:06.400102Z",
+     "shell.execute_reply": "2026-10-14T09:55:06.399004Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "    u     r   theory       MC     z\n",
+      "0.100   1.0   0.1000   0.0942  1.72\n",
+      "0.100   5.0   0.3571   0.3589  0.38\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0.300   1.0   0.3000   0.3014  0.28\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0.300   5.0   0.6818   0.6830  0.34\n",
+      "0.477   5.0   0.8202   0.8198  0.13\n",
+      "0.477  20.0   0.9480   0.9480  0.07\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0.800   3.0   0.9231   0.9221  0.52\n",
+      "worst |z| = 1.72\n",
+      "\n",
+      "anchors (from the papers; see the note): VERINA best model: spec sound-and-complete 52.3% => u ~ 0.477; CLEVER: spec generation + equivalence proof solved <= 1/161 by few-shot LLMs; miniF2F-Lean revisited: formal/informal mismatch in > 1/2 of 488 problems.\n"
+     ]
+    }
+   ],
+   "source": [
+    "def p_wrong_given_accepted(u, r): return u * r / (u * r + 1 - u)\n",
+    "rng = random.Random(3); worst = 0.0\n",
+    "print(f'{\"u\":>5s} {\"r\":>5s} {\"theory\":>8s} {\"MC\":>8s} {\"z\":>5s}')\n",
+    "for u, r_ in [(0.1, 1), (0.1, 5), (0.3, 1), (0.3, 5), (0.477, 5), (0.477, 20), (0.8, 3)]:\n",
+    "    a_c = 0.04; a_w = min(1.0, a_c * r_)\n",
+    "    acc = wrong = 0\n",
+    "    for _ in range(200000):\n",
+    "        w_ = rng.random() < u\n",
+    "        if rng.random() < (a_w if w_ else a_c):\n",
+    "            acc += 1; wrong += w_\n",
+    "    mc = wrong / acc; th = p_wrong_given_accepted(u, a_w / a_c)\n",
+    "    se = (th * (1 - th) / acc) ** 0.5; z = abs(mc - th) / se; worst = max(worst, z)\n",
+    "    print(f'{u:5.3f} {a_w/a_c:5.1f} {th:8.4f} {mc:8.4f} {z:5.2f}')\n",
+    "assert worst < 4\n",
+    "print('worst |z| =', round(worst, 2))\n",
+    "print('\\nanchors (from the papers; see the note): VERINA best model: spec sound-and-complete 52.3% => u ~ 0.477; '\n",
+    "      'CLEVER: spec generation + equivalence proof solved <= 1/161 by few-shot LLMs; miniF2F-Lean revisited: formal/informal mismatch in > 1/2 of 488 problems.')\n",
+    "L.close()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

code/notebooks/0306-benchmarks-landscape.ipynb

@@ -0,0 +1,479 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "3a638e03",
+   "metadata": {},
+   "source": [
+    "# 0306 - Verified-code benchmarks: what each measures, and how to read the numbers\n",
+    "\n",
+    "Scratch-pad for note `0306 - Verified Code Benchmarks.pdf`. The benchmarks (VERINA, CLEVER, FVAPPS, DafnyBench, miniCodeProps) differ in\n",
+    "*task*, *who writes the spec*, *how the spec is validated* and *what the checker trusts*; that is the table in the note. This notebook\n",
+    "verifies the **arithmetic** the note depends on (naive first, efficient second, assert agreement):\n",
+    "\n",
+    "1. **pass@k** (Chen et al. 2021, eq. (1), Fig. 3): NAIVE exact combinatorics; EFFICIENT stable product. Unbiasedness of the estimator\n",
+    "   (exact sum over the binomial), and the bias of the plug-in `1-(1-p_hat)^k`.\n",
+    "2. **Small benchmarks**: Wilson vs Wald interval, coverage by Monte Carlo, applied to the headline numbers of each paper.\n",
+    "3. **Composite tasks**: Frechet bounds for `P(code and spec and proof)`; checked on random joint distributions.\n",
+    "4. **Consistency checks** of numbers quoted from the papers (sums, ratios) and one *derived* baseline (DafnyBench's no-hint rate).\n",
+    "\n",
+    "No GPU, no Lean in this notebook (pure arithmetic)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "542867e2",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:53:51.893124Z",
+     "iopub.status.busy": "2026-10-14T09:53:51.892857Z",
+     "iopub.status.idle": "2026-10-14T09:53:52.032196Z",
+     "shell.execute_reply": "2026-10-14T09:53:52.030551Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
+    "import math, random, itertools\n",
+    "from fractions import Fraction\n",
+    "import numpy as np\n",
+    "random.seed(306); np.random.seed(306)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3ae94319",
+   "metadata": {},
+   "source": [
+    "## 1. pass@k\n",
+    "Given `n` samples of which `c` are correct, a random `k`-subset contains no correct sample with probability `C(n-c,k)/C(n,k)`,\n",
+    "so `pass@k = 1 - C(n-c,k)/C(n,k)`. Stable form (cancel the factorials):\n",
+    "`C(n-c,k)/C(n,k) = prod_{i=n-c+1}^{n} (1 - k/i)`  (for `n-c >= k`; else the ratio is 0)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "806c53e0",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:53:52.034678Z",
+     "iopub.status.busy": "2026-10-14T09:53:52.034346Z",
+     "iopub.status.idle": "2026-10-14T09:53:52.050351Z",
+     "shell.execute_reply": "2026-10-14T09:53:52.049783Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "stable product vs exact rational on 303 (n,c,k) triples: max abs error 6.11e-16\n",
+      "naive float64 factorial ratio at n=200 (170! < 1.8e308 < 171!): nan  (exact: 0.660225623763214 )\n",
+      "naive float is fine for small n: True\n"
+     ]
+    }
+   ],
+   "source": [
+    "def pass_at_k_exact(n, c, k):                       # NAIVE: exact rational arithmetic from the definition\n",
+    "    if n - c < k: return Fraction(1)\n",
+    "    return 1 - Fraction(math.comb(n - c, k), math.comb(n, k))\n",
+    "def _ffact(m):                                        # factorial accumulated in float64 (what a vectorised implementation does)\n",
+    "    return float(np.prod(np.arange(1, m + 1, dtype=np.float64)))\n",
+    "def pass_at_k_naive_float(n, c, k):                  # NAIVE in floats: the textbook factorial ratio C(n-c,k)/C(n,k)\n",
+    "    if n - c < k: return 1.0\n",
+    "    with np.errstate(all='ignore'):\n",
+    "        cnk = lambda a, b: _ffact(a) / (_ffact(b) * _ffact(a - b))\n",
+    "        return 1.0 - cnk(n - c, k) / cnk(n, k)\n",
+    "def pass_at_k_stable(n, c, k):                       # EFFICIENT: Chen et al., Fig. 3\n",
+    "    if n - c < k: return 1.0\n",
+    "    return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1))\n",
+    "\n",
+    "worst = 0.0; count = 0\n",
+    "for n in (5, 10, 20, 50, 100, 170, 200, 500):\n",
+    "    for c in sorted({0, 1, 2, n // 10, n // 3, n // 2, n - 1, n}):\n",
+    "        for k in (1, 2, 5, 10, 50, 100):\n",
+    "            if k > n: continue\n",
+    "            ex = float(pass_at_k_exact(n, c, k)); st = pass_at_k_stable(n, c, k)\n",
+    "            worst = max(worst, abs(ex - st)); count += 1\n",
+    "assert worst < 1e-12\n",
+    "print(f'stable product vs exact rational on {count} (n,c,k) triples: max abs error {worst:.2e}')\n",
+    "v200 = pass_at_k_naive_float(200, 20, 10)\n",
+    "print('naive float64 factorial ratio at n=200 (170! < 1.8e308 < 171!):', v200, ' (exact:', float(pass_at_k_exact(200, 20, 10)), ')')\n",
+    "assert not math.isfinite(v200)\n",
+    "print('naive float is fine for small n:', abs(pass_at_k_naive_float(20, 5, 3) - float(pass_at_k_exact(20, 5, 3))) < 1e-12)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "46406668",
+   "metadata": {},
+   "source": [
+    "### Unbiasedness (exact) and the bias of the plug-in estimator\n",
+    "For a task with true single-sample success probability `p`, `c ~ Binomial(n,p)`. Claim: `E[pass_at_k(n,c,k)] = 1-(1-p)^k`.\n",
+    "Proof sketch: the random `k`-subset of the `n` samples is itself `k` i.i.d. draws, so it contains no success with probability `(1-p)^k`.\n",
+    "Check by exact summation. The plug-in `1-(1-c/n)^k` is biased (Chen et al., Appendix A)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "a46d5fa0",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:53:52.053085Z",
+     "iopub.status.busy": "2026-10-14T09:53:52.052843Z",
+     "iopub.status.idle": "2026-10-14T09:53:52.063079Z",
+     "shell.execute_reply": "2026-10-14T09:53:52.061837Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "  n     p   k  truth 1-(1-p)^k  E[unbiased]     |err|  E[plug-in]      bias\n",
+      " 10  0.10   5     0.4095100000 0.4095100000   2.2e-16  0.34851642  -0.06099\n",
+      " 10  0.30   5     0.8319300000 0.8319300000   6.7e-16  0.75974023  -0.07219\n",
+      " 20  0.05  10     0.4012630608 0.4012630608   6.1e-16  0.33649458  -0.06477\n",
+      " 20  0.50   3     0.8750000000 0.8750000000   0.0e+00  0.85625000  -0.01875\n",
+      " 50  0.02  20     0.3323920282 0.3323920282   6.7e-16  0.28490846  -0.04748\n",
+      "100  0.01  50     0.3949939329 0.3949939329   6.7e-16  0.32746853  -0.06753\n",
+      "unbiasedness: max |E[estimator] - truth| = 6.7e-16;  the plug-in estimator is biased (negative here: it under-estimates pass@k).\n"
+     ]
+    }
+   ],
+   "source": [
+    "def expect_estimators(n, p, k):\n",
+    "    e_unb = e_plug = 0.0\n",
+    "    for c in range(n + 1):\n",
+    "        w = math.comb(n, c) * p ** c * (1 - p) ** (n - c)\n",
+    "        e_unb += w * pass_at_k_stable(n, c, k)\n",
+    "        e_plug += w * (1 - (1 - c / n) ** k)\n",
+    "    return e_unb, e_plug\n",
+    "print(f'{\"n\":>3s} {\"p\":>5s} {\"k\":>3s} {\"truth 1-(1-p)^k\":>16s} {\"E[unbiased]\":>12s} {\"|err|\":>9s} {\"E[plug-in]\":>11s} {\"bias\":>9s}')\n",
+    "max_err = 0.0\n",
+    "for n, p, k in [(10, 0.1, 5), (10, 0.3, 5), (20, 0.05, 10), (20, 0.5, 3), (50, 0.02, 20), (100, 0.01, 50)]:\n",
+    "    truth = 1 - (1 - p) ** k\n",
+    "    u, pl = expect_estimators(n, p, k)\n",
+    "    max_err = max(max_err, abs(u - truth))\n",
+    "    print(f'{n:3d} {p:5.2f} {k:3d} {truth:16.10f} {u:12.10f} {abs(u-truth):9.1e} {pl:11.8f} {pl-truth:9.5f}')\n",
+    "assert max_err < 1e-12\n",
+    "print(f'unbiasedness: max |E[estimator] - truth| = {max_err:.1e};  the plug-in estimator is biased (negative here: it under-estimates pass@k).')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3aa2cde7",
+   "metadata": {},
+   "source": [
+    "## 2. Small benchmarks need intervals\n",
+    "VERINA 189 tasks, CLEVER 161, miniCodeProps 201, FVAPPS baseline 100 samples / 406 theorems, DafnyBench 782. Treat each task as a Bernoulli trial\n",
+    "(NAMED APPROXIMATION: tasks exchangeable and independent; FVAPPS theorems within a sample are correlated, so its interval is too narrow).\n",
+    "Wald: `p_hat +- 1.96 sqrt(p_hat(1-p_hat)/n)`. Wilson: solve `|p_hat - p| = z sqrt(p(1-p)/n)` for `p`:\n",
+    "`centre = (p_hat + z^2/2n)/(1 + z^2/n)`, `half = z sqrt(p_hat(1-p_hat)/n + z^2/4n^2)/(1 + z^2/n)`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "dfbe9e5a",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:53:52.065104Z",
+     "iopub.status.busy": "2026-10-14T09:53:52.064886Z",
+     "iopub.status.idle": "2026-10-14T09:53:52.480161Z",
+     "shell.execute_reply": "2026-10-14T09:53:52.479404Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Wilson endpoints solve (p_hat - p)^2 = z^2 p(1-p)/n exactly (residual < 1e-12) for 5 test cases\n",
+      "   n      p  Wald cover  Wilson cover\n",
+      " 161  0.005      0.5518        0.9539\n",
+      " 161  0.030      0.8595        0.9675\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      " 189  0.049      0.9451        0.9601\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      " 201  0.350      0.9536        0.9462\n",
+      " 406  0.300      0.9497        0.9498\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      " 782  0.680      0.9496        0.9498\n",
+      "\n",
+      "benchmark      setting                                                                  rate  Wilson 95% interval\n",
+      "CLEVER         best few-shot LLM, end-to-end                                            0.6%  [ 0.1%,  3.4%]\n",
+      "VERINA         best model proof success, pass@1 (4.9% of 189 ~ 9.3 tasks; k rounded)    4.8%  [ 2.5%,  8.8%]\n",
+      "VERINA         best model spec sound+complete (52.3% of 189 ~ 98.9; k rounded)         52.4%  [45.3%, 59.4%]\n",
+      "miniCodeProps  GPT-4o pass@32 + refinement, overall                                    37.3%  [30.9%, 44.2%]\n",
+      "miniCodeProps  GPT-4o pass@32 + refinement, medium+hard                                 7.0%  [ 3.6%, 13.1%]\n",
+      "FVAPPS         Claude 3.5 Sonnet, theorems of 100 sampled problems                     29.8%  [25.6%, 34.4%]\n",
+      "FVAPPS         Gemini 1.5 Pro                                                          18.2%  [14.8%, 22.3%]\n",
+      "DafnyBench     best model, Table 3: 67.8% of 782 = 530                                 67.8%  [64.4%, 71.0%]\n",
+      "DafnyBench     no-LLM baseline, Table 3: 26.9% of 782 ~ 210                            26.9%  [23.9%, 30.1%]\n"
+     ]
+    }
+   ],
+   "source": [
+    "Z = 1.959963984540054\n",
+    "def wald(k, n):\n",
+    "    p = k / n; h = Z * math.sqrt(p * (1 - p) / n); return max(0, p - h), min(1, p + h)\n",
+    "def wilson(k, n):\n",
+    "    p = k / n; d = 1 + Z * Z / n\n",
+    "    c = (p + Z * Z / (2 * n)) / d; h = Z * math.sqrt(p * (1 - p) / n + Z * Z / (4 * n * n)) / d\n",
+    "    return c - h, c + h\n",
+    "\n",
+    "# derivation check: the Wilson endpoints are exactly the roots of (p_hat - p)^2 = z^2 p (1-p)/n\n",
+    "for k, n in [(1, 161), (9, 189), (70, 201), (121, 406), (530, 782)]:\n",
+    "    lo, hi = wilson(k, n)\n",
+    "    for p_ in (lo, hi):\n",
+    "        assert abs((k / n - p_) ** 2 - Z * Z * p_ * (1 - p_) / n) < 1e-12\n",
+    "print('Wilson endpoints solve (p_hat - p)^2 = z^2 p(1-p)/n exactly (residual < 1e-12) for 5 test cases')\n",
+    "\n",
+    "rng = np.random.default_rng(1)\n",
+    "def coverage(n, p, reps=40000):\n",
+    "    ks = rng.binomial(n, p, size=reps)\n",
+    "    cov_w = cov_l = 0\n",
+    "    for k in ks:\n",
+    "        lo, hi = wald(int(k), n); cov_w += (lo <= p <= hi)\n",
+    "        lo, hi = wilson(int(k), n); cov_l += (lo <= p <= hi)\n",
+    "    return cov_w / reps, cov_l / reps\n",
+    "print(f'{\"n\":>4s} {\"p\":>6s} {\"Wald cover\":>11s} {\"Wilson cover\":>13s}')\n",
+    "cov_table = {}\n",
+    "for n, p in [(161, 0.005), (161, 0.03), (189, 0.049), (201, 0.35), (406, 0.30), (782, 0.68)]:\n",
+    "    cw, cl = coverage(n, p); cov_table[(n, p)] = (cw, cl)\n",
+    "    print(f'{n:4d} {p:6.3f} {cw:11.4f} {cl:13.4f}')\n",
+    "assert cov_table[(161, 0.005)][0] < 0.6 and cov_table[(161, 0.005)][1] > 0.90        # Wald collapses at tiny p; Wilson does not\n",
+    "assert all(cl > 0.92 for (_, cl) in cov_table.values())\n",
+    "\n",
+    "print()\n",
+    "HEADLINES = [   # (benchmark, what, successes k, n)  -- k reconstructed from the reported rate and n where the paper gives a rate\n",
+    " ('CLEVER', 'best few-shot LLM, end-to-end', 1, 161),\n",
+    " ('VERINA', 'best model proof success, pass@1 (4.9% of 189 ~ 9.3 tasks; k rounded)', 9, 189),\n",
+    " ('VERINA', 'best model spec sound+complete (52.3% of 189 ~ 98.9; k rounded)', 99, 189),\n",
+    " ('miniCodeProps', 'GPT-4o pass@32 + refinement, overall', 75, 201),\n",
+    " ('miniCodeProps', 'GPT-4o pass@32 + refinement, medium+hard', 8, 115),\n",
+    " ('FVAPPS', 'Claude 3.5 Sonnet, theorems of 100 sampled problems', 121, 406),\n",
+    " ('FVAPPS', 'Gemini 1.5 Pro', 74, 406),\n",
+    " ('DafnyBench', 'best model, Table 3: 67.8% of 782 = 530', 530, 782),\n",
+    " ('DafnyBench', 'no-LLM baseline, Table 3: 26.9% of 782 ~ 210', 210, 782),\n",
+    "]\n",
+    "print(f'{\"benchmark\":14s} {\"setting\":70s} {\"rate\":>6s}  Wilson 95% interval')\n",
+    "for b, what, k, n in HEADLINES:\n",
+    "    lo, hi = wilson(k, n)\n",
+    "    print(f'{b:14s} {what:70s} {100*k/n:5.1f}%  [{100*lo:4.1f}%, {100*hi:4.1f}%]')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bc025e01",
+   "metadata": {},
+   "source": [
+    "## 3. Composite tasks: Frechet bounds\n",
+    "End-to-end success `A and B and C` (code, spec, proof; or CLEVER's `QED1 and QED2`) is bounded by its stages whatever the dependence:\n",
+    "`max(0, sum P(.) - (m-1)) <= P(all) <= min P(.)`. Check on random joint distributions of `m` binary stage outcomes."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "8b0beb96",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:53:52.481823Z",
+     "iopub.status.busy": "2026-10-14T09:53:52.481664Z",
+     "iopub.status.idle": "2026-10-14T09:53:53.003946Z",
+     "shell.execute_reply": "2026-10-14T09:53:53.003015Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "20000 random joint distributions (m = 2 or 3 stages): Frechet bound violations = 0\n",
+      "VERINA stage marginals (code 72.6%, spec 52.3%, proof 4.9%) =>  P(all three) in [0.000, 0.049]\n",
+      "comonotone coupling attains the upper bound: 0.0491 vs min = 0.049\n",
+      "independent stages would give the product 0.01861 (inside the bounds; real tasks are positively dependent)\n",
+      "lower Frechet bound attained by an explicit coupling on 6 marginals (exact rational arithmetic, grid of 1/20; e.g. (0.7,0.6,0.9) -> 0.2, (0.5,0.5,0.5) -> 0)\n"
+     ]
+    }
+   ],
+   "source": [
+    "def frechet(ps):\n",
+    "    return max(0.0, sum(ps) - (len(ps) - 1)), min(ps)\n",
+    "viol = 0; tight_lo = tight_hi = 0\n",
+    "for trial in range(20000):\n",
+    "    m = random.choice([2, 3])\n",
+    "    w = np.random.dirichlet(np.ones(2 ** m) * random.choice([0.2, 1.0, 5.0]))\n",
+    "    outcomes = list(itertools.product([0, 1], repeat=m))\n",
+    "    ps = [sum(w[i] for i, o in enumerate(outcomes) if o[j]) for j in range(m)]\n",
+    "    joint = sum(w[i] for i, o in enumerate(outcomes) if all(o))\n",
+    "    lo, hi = frechet(ps)\n",
+    "    viol += not (lo - 1e-12 <= joint <= hi + 1e-12)\n",
+    "assert viol == 0\n",
+    "print('20000 random joint distributions (m = 2 or 3 stages): Frechet bound violations =', viol)\n",
+    "# tightness: comonotone coupling attains the upper bound, counter-monotone (m=2) the lower\n",
+    "ps = (0.726, 0.523, 0.049)\n",
+    "lo_, hi_ = frechet(ps)\n",
+    "print(f'VERINA stage marginals (code 72.6%, spec 52.3%, proof 4.9%) =>  P(all three) in [{lo_:.3f}, {hi_:.3f}]')\n",
+    "assert frechet(ps) == (0.0, 0.049)\n",
+    "# attained: draw U~U(0,1); stage j succeeds iff U < p_j  => joint = min p\n",
+    "U = np.random.rand(2_000_000); joint_co = np.mean((U < ps[0]) & (U < ps[1]) & (U < ps[2]))\n",
+    "print(f'comonotone coupling attains the upper bound: {joint_co:.4f} vs min = {min(ps)}')\n",
+    "assert abs(joint_co - min(ps)) < 2e-3\n",
+    "print('independent stages would give the product', round(math.prod(ps), 5), '(inside the bounds; real tasks are positively dependent)')\n",
+    "\n",
+    "# the LOWER bound is attained too: put the complements A_j^c on consecutive intervals of the circle [0,1) (disjoint while sum(1-p) <= 1, covering otherwise)\n",
+    "from fractions import Fraction as Fr\n",
+    "def lower_attained(ps, cells=20):\n",
+    "    ps = [Fr(p).limit_denominator(cells) for p in ps]\n",
+    "    comp = [int((1 - p) * cells) for p in ps]                 # |A_j^c| in cells of the grid\n",
+    "    pos = 0; ok = [True] * cells\n",
+    "    for c in comp:\n",
+    "        for t in range(c): ok[(pos + t) % cells] = False       # cells where stage j FAILS\n",
+    "        pos += c\n",
+    "    joint = Fr(sum(ok), cells)\n",
+    "    return joint, max(Fr(0), sum(ps) - (len(ps) - 1))\n",
+    "for ps_ in [(0.7, 0.6, 0.9), (0.5, 0.5, 0.5), (0.95, 0.9), (0.3, 0.4), (0.8, 0.8, 0.8, 0.8), (0.75, 0.5, 0.05)]:\n",
+    "    j_, b_ = lower_attained(ps_)\n",
+    "    assert j_ == b_, (ps_, j_, b_)\n",
+    "print('lower Frechet bound attained by an explicit coupling on 6 marginals (exact rational arithmetic, grid of 1/20; e.g. (0.7,0.6,0.9) -> 0.2, (0.5,0.5,0.5) -> 0)')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b98fb644",
+   "metadata": {},
+   "source": [
+    "## 4. Consistency checks of quoted numbers (and one derived baseline)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "291ff898",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-10-14T09:53:53.005521Z",
+     "iopub.status.busy": "2026-10-14T09:53:53.005355Z",
+     "iopub.status.idle": "2026-10-14T09:53:53.013532Z",
+     "shell.execute_reply": "2026-10-14T09:53:53.012678Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "OK   DafnyBench sources sum to 782 (556 GitHub + 62 Clover + 164 dafny-synthesis)\n",
+      "OK   DafnyBench section 3.1 no-hint count vs Table 3 \"No LLM\" baseline (113+23+72 = 208 of 782 = 26.6%; Table 3 states the no-op baseline itself: 26.9% = 210 programs, 2 more than the section-3.1 count)\n",
+      "OK   DafnyBench Table 3 +- for GPT-3.5 Turbo (binomial s.e. 1.78 vs reported 1.8)\n",
+      "OK   DafnyBench Table 3 +- for GPT-4 Turbo (binomial s.e. 1.75 vs reported 1.8)\n",
+      "OK   DafnyBench Table 3 +- for Claude 3 Opus (binomial s.e. 1.67 vs reported 1.7)\n",
+      "OK   DafnyBench Table 3 +- for CodeLlama-7b (binomial s.e. 1.61 vs reported 1.6)\n",
+      "OK   miniCodeProps GPT-4o medley (65/86 = 75.58% vs reported 75.6%)\n",
+      "OK   miniCodeProps GPT-4o med+hard (5/115 = 4.35% vs reported 4.34%)\n",
+      "OK   miniCodeProps GPT-4o overall (70/201 = 34.83% vs reported 34.8%)\n",
+      "OK   miniCodeProps refine medley (67/86 = 77.91% vs reported 77.9%)\n",
+      "OK   miniCodeProps refine med+hard (8/115 = 6.96% vs reported 6.96%)\n",
+      "OK   miniCodeProps refine overall (75/201 = 37.31% vs reported 37.3%)\n",
+      "OK   miniCodeProps ntp medley (62/86 = 72.09% vs reported 72.1%)\n",
+      "OK   miniCodeProps ntp med+hard (10/115 = 8.70% vs reported 8.69%)\n",
+      "OK   miniCodeProps ntp overall (72/201 = 35.82% vs reported 35.8%)\n",
+      "OK   miniCodeProps 86 + 115 = 201 \n",
+      "OK   miniCodeProps overall = sum of splits (GPT-4o) \n",
+      "OK   FVAPPS Sonnet 121/406 ~ 30% (29.80%)\n",
+      "FLAG FVAPPS Gemini 74/406 vs paper text \"18.5%\" (18.23%; abstract says 18%) -> small internal inconsistency in the paper\n",
+      "OK   FVAPPS tiers nested: 1083 <= 2089 <= 4715 (23.0% pass all five stages; 44.3% guarded)\n",
+      "OK   CLEVER 161 of 164 HumanEval \n",
+      "OK   VERINA ground-truth proofs 46 of 189 (24.3% of tasks have a reference proof)\n",
+      "\n",
+      "22 checks, 1 flagged (the FVAPPS 18.5% vs 74/406 = 18.2% discrepancy is the only one expected)\n"
+     ]
+    }
+   ],
+   "source": [
+    "checks = []\n",
+    "def chk(name, ok, detail=''):\n",
+    "    checks.append((name, ok, detail)); print(f'{\"OK  \" if ok else \"FLAG\"} {name} {detail}')\n",
+    "# DafnyBench (arXiv 2406.08467 section 3.1): sources of the 782 ground-truth programs\n",
+    "chk('DafnyBench sources sum to 782', 556 + 62 + 164 == 782, '(556 GitHub + 62 Clover + 164 dafny-synthesis)')\n",
+    "no_hint = 113 + 23 + 72\n",
+    "chk('DafnyBench section 3.1 no-hint count vs Table 3 \"No LLM\" baseline', abs(100 * no_hint / 782 - 26.9) < 0.5,\n",
+    "    f'(113+23+72 = {no_hint} of 782 = {100*no_hint/782:.1f}%; Table 3 states the no-op baseline itself: 26.9% = {round(0.269*782)} programs, {round(0.269*782)-no_hint} more than the section-3.1 count)')\n",
+    "# Table 3 reports '+-' = one binomial standard error at N = 782\n",
+    "for name, p_, se_rep in [('GPT-3.5 Turbo', 0.440, 1.8), ('GPT-4 Turbo', 0.598, 1.8), ('Claude 3 Opus', 0.678, 1.7), ('CodeLlama-7b', 0.280, 1.6)]:\n",
+    "    chk(f'DafnyBench Table 3 +- for {name}', abs(100 * math.sqrt(p_ * (1 - p_) / 782) - se_rep) < 0.06, f'(binomial s.e. {100*math.sqrt(p_*(1-p_)/782):.2f} vs reported {se_rep})')\n",
+    "# miniCodeProps (arXiv 2406.11915, Table 3)\n",
+    "for name, k, n, rep in [('GPT-4o medley', 65, 86, 75.6), ('GPT-4o med+hard', 5, 115, 4.34), ('GPT-4o overall', 70, 201, 34.8),\n",
+    "                        ('refine medley', 67, 86, 77.9), ('refine med+hard', 8, 115, 6.96), ('refine overall', 75, 201, 37.3),\n",
+    "                        ('ntp medley', 62, 86, 72.1), ('ntp med+hard', 10, 115, 8.69), ('ntp overall', 72, 201, 35.8)]:\n",
+    "    chk(f'miniCodeProps {name}', abs(100 * k / n - rep) < 0.06, f'({k}/{n} = {100*k/n:.2f}% vs reported {rep}%)')\n",
+    "chk('miniCodeProps 86 + 115 = 201', 86 + 115 == 201)\n",
+    "chk('miniCodeProps overall = sum of splits (GPT-4o)', 65 + 5 == 70 and 67 + 8 == 75 and 62 + 10 == 72)\n",
+    "# FVAPPS (arXiv 2502.05714 section III)\n",
+    "chk('FVAPPS Sonnet 121/406 ~ 30%', abs(100 * 121 / 406 - 30) < 0.5, f'({100*121/406:.2f}%)')\n",
+    "chk('FVAPPS Gemini 74/406 vs paper text \"18.5%\"', abs(100 * 74 / 406 - 18.5) < 0.1, f'({100*74/406:.2f}%; abstract says 18%) -> small internal inconsistency in the paper')\n",
+    "chk('FVAPPS tiers nested: 1083 <= 2089 <= 4715', 1083 <= 2089 <= 4715, f'({100*1083/4715:.1f}% pass all five stages; {100*2089/4715:.1f}% guarded)')\n",
+    "# CLEVER (arXiv 2505.13938)\n",
+    "chk('CLEVER 161 of 164 HumanEval', 161 <= 164)\n",
+    "# VERINA (arXiv 2505.23135): test-suite medians and the relation count\n",
+    "chk('VERINA ground-truth proofs 46 of 189', abs(100 * 46 / 189 - 24.3) < 0.1, f'({100*46/189:.1f}% of tasks have a reference proof)')\n",
+    "n_fail = sum(not ok for _, ok, _ in checks)\n",
+    "print(f'\\n{len(checks)} checks, {n_fail} flagged (the FVAPPS 18.5% vs 74/406 = 18.2% discrepancy is the only one expected)')\n",
+    "assert n_fail == 1"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8e17d222",
+   "metadata": {},
+   "source": [
+    "## What the table in the note is built from\n",
+    "The benchmark comparison table (task / who writes the spec / spec validation / checker / headline) is hand-assembled from each paper's own text in the\n",
+    "note; the numbers above are the ones that can be recomputed. Reading rule: **a headline number needs (a) the task definition, (b) a trivial baseline, (c) an interval.**"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

code/notebooks/0307-aeneas-semantics.ipynb

@@ -0,0 +1,438 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 0307 - Aeneas semantics: Rust (safe subset) -> LLBC -> pure Lean, the Result monad, and what a proof of `not P` certifies\n",
+    "Companion to `notes/0307 - Aeneas Semantics.pdf`.  Pure Python, exhaustive enumeration over `u8`; no Lean needed.\n",
+    "The \"Rust abstract machine\" below is a small stateful interpreter (store, borrows, panics, fuel).  The \"translation\" is a pure function\n",
+    "into `Result` with explicit threading of borrowed values (backward functions).  Every check compares them on the **whole** finite domain.\n",
+    "This is a finite test of the refinement statement, not a proof of Aeneas' soundness theorem (that is a meta-theorem about the tool, see the note)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "source": [
+    "import os\n",
+    "os.environ[\"OMP_NUM_THREADS\"] = \"1\"\n",
+    "import itertools, math\n",
+    "from dataclasses import dataclass\n",
+    "CHECKS = []\n",
+    "def check(name, ok, detail=\"\"):\n",
+    "    CHECKS.append((name, bool(ok), detail))\n",
+    "    assert ok, (name, detail)\n",
+    "    print(f\"[ok] {name}  {detail}\")"
+   ],
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. The Result monad: ok / fail(e) / div\n",
+    "`Result a = ok a | fail Error | div`.  `bind` propagates `fail` and `div`.  Errors we use: `integerOverflow`, `arrayOutOfBounds`, `panic`.\n",
+    "Monad laws are checked on a finite carrier (all values of `u8`, three error tags, `div`)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "source": [
+    "@dataclass(frozen=True)\n",
+    "class R:\n",
+    "    tag: str            # 'ok' | 'fail' | 'div'\n",
+    "    v: object = None    # value for ok, error name for fail\n",
+    "def ok(v): return R('ok', v)\n",
+    "def fail(e): return R('fail', e)\n",
+    "DIV = R('div')\n",
+    "def bind(m, f):\n",
+    "    return f(m.v) if m.tag == 'ok' else m\n",
+    "U8 = range(256)\n",
+    "carrier = [ok(v) for v in U8] + [fail('integerOverflow'), fail('arrayOutOfBounds'), fail('panic'), DIV]\n",
+    "f = lambda x: ok(x + 1) if x < 255 else fail('integerOverflow')   # a fallible step\n",
+    "g = lambda x: ok(x // 2) if x % 3 else fail('panic')\n",
+    "left  = all(bind(ok(v), f) == f(v) for v in U8)\n",
+    "right = all(bind(m, ok) == m for m in carrier)\n",
+    "assoc = all(bind(bind(m, f), g) == bind(m, lambda x: bind(f(x), g)) for m in carrier)\n",
+    "check(\"1. monad laws on 260 values: left identity, right identity, associativity\", left and right and assoc)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "[ok] 1. monad laws on 260 values: left identity, right identity, associativity  \n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Naive translation fails: integers without overflow\n",
+    "Rust (debug / `overflow-checks`): `fn avg(a: u8, b: u8) -> u8 { (a + b) / 2 }`.  The abstract machine panics iff `a + b > 255`.\n",
+    "NAIVE translation: `a, b` as mathematical integers, no failure case.  CORRECT translation: `let s <- a + b; s / 2` with checked `+` in `Result`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "source": [
+    "def rust_avg(a, b):                       # abstract machine, checked arithmetic\n",
+    "    s = a + b\n",
+    "    if s > 255: return ('panic', 'integerOverflow')\n",
+    "    return ('ret', s // 2)\n",
+    "def tr_naive_avg(a, b):  return ok((a + b) // 2)            # no Result\n",
+    "def u8_add(a, b):  return ok(a + b) if a + b <= 255 else fail('integerOverflow')\n",
+    "def tr_avg(a, b):  return bind(u8_add(a, b), lambda s: ok(s // 2))\n",
+    "def obs(r):        return ('ret', r.v) if r.tag == 'ok' else (('panic', r.v) if r.tag == 'fail' else ('div', None))\n",
+    "DOM = [(a, b) for a in U8 for b in U8]\n",
+    "bad_naive = sum(1 for a, b in DOM if obs(tr_naive_avg(a, b)) != rust_avg(a, b))\n",
+    "bad_good  = sum(1 for a, b in DOM if obs(tr_avg(a, b)) != rust_avg(a, b))\n",
+    "panics = sum(1 for a, b in DOM if rust_avg(a, b)[0] == 'panic')\n",
+    "check(\"2. naive translation disagrees with the machine on exactly the panicking inputs\", bad_naive == panics == 255 * 256 // 2, f\"disagree={bad_naive}, panics={panics}, 255*256/2={255*256//2}\")\n",
+    "check(\"2. Result translation agrees with the machine on all 65536 inputs\", bad_good == 0, f\"disagree={bad_good}\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "[ok] 2. naive translation disagrees with the machine on exactly the panicking inputs  disagree=32640, panics=32640, 255*256/2=32640\n",
+      "[ok] 2. Result translation agrees with the machine on all 65536 inputs  disagree=0\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Borrow elimination: `&mut` becomes \"return the new value\" (forward) and, when a borrow is returned, a backward function\n",
+    "Machine: a store maps addresses to values; `&mut x` is an address.  `fn bump(x: &mut u8) { *x += 1 }`.\n",
+    "Translation: `bump (x : U8) : Result U8` returns the new value and the caller rebinds."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "source": [
+    "def machine_bump(store, addr):            # returns (outcome, store')\n",
+    "    v = store[addr] + 1\n",
+    "    if v > 255: return ('panic', 'integerOverflow'), store          # store unchanged when the panic happens before the write\n",
+    "    s2 = dict(store); s2[addr] = v\n",
+    "    return ('ret', None), s2\n",
+    "def tr_bump(x): return u8_add(x, 1)\n",
+    "good = 0\n",
+    "for x in U8:\n",
+    "    out, st = machine_bump({0: x}, 0)\n",
+    "    r = tr_bump(x)\n",
+    "    if out[0] == 'ret': good += (r.tag == 'ok' and r.v == st[0])\n",
+    "    else: good += (r.tag == 'fail' and r.v == out[1])\n",
+    "check(\"3a. bump: machine (store, outcome) = translation (returned value, Result) for all u8\", good == 256)\n",
+    "# Returned mutable borrow:  fn get_mut(v: &mut Vec<u8>, i: usize) -> &mut u8 ;  caller:  *get_mut(&mut v, i) += k\n",
+    "def machine_incr_at(vec, i, k):\n",
+    "    heap = {j: x for j, x in enumerate(vec)}            # elements are addressable places\n",
+    "    if i >= len(vec): return ('panic', 'arrayOutOfBounds'), vec\n",
+    "    v = heap[i] + k\n",
+    "    if v > 255: return ('panic', 'integerOverflow'), vec\n",
+    "    heap[i] = v\n",
+    "    return ('ret', None), [heap[j] for j in range(len(vec))]\n",
+    "def tr_get_mut(v, i):                      # forward value and BACKWARD function  back : u8 -> Vec\n",
+    "    if i >= len(v): return fail('arrayOutOfBounds')\n",
+    "    back = lambda new, v=v, i=i: tuple(new if j == i else x for j, x in enumerate(v))\n",
+    "    return ok((v[i], back))\n",
+    "def tr_incr_at(v, i, k):\n",
+    "    return bind(tr_get_mut(v, i), lambda p: bind(u8_add(p[0], k), lambda new: ok(p[1](new))))\n",
+    "cases = ties = 0\n",
+    "for n in range(0, 4):\n",
+    "    for vec in itertools.product([0, 1, 254, 255], repeat=n):\n",
+    "        for i in range(0, 5):\n",
+    "            for k in (0, 1, 2):\n",
+    "                out, st = machine_incr_at(list(vec), i, k)\n",
+    "                r = tr_incr_at(vec, i, k)\n",
+    "                cases += 1\n",
+    "                if out[0] == 'ret': ties += (r.tag == 'ok' and list(r.v) == st)\n",
+    "                else: ties += (r.tag == 'fail' and r.v == out[1])\n",
+    "check(\"3b. returned &mut: machine vs forward value + backward function, vec len<=3, i<=4, k<=2\", ties == cases, f\"{ties}/{cases} cases\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "[ok] 3a. bump: machine (store, outcome) = translation (returned value, Result) for all u8  \n",
+      "[ok] 3b. returned &mut: machine vs forward value + backward function, vec len<=3, i<=4, k<=2  1275/1275 cases\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Loops become recursion; non-termination is `div`\n",
+    "`fn spin(x: u8) { while x != 0 { } }`.  The machine is run with a fuel bound F; `div` is what survives every F.\n",
+    "`fn sum_to(n: u8) -> u8 { let mut s = 0; let mut i = 0; while i < n { i += 1; s += i } s }` : the accumulator overflows once n(n+1)/2 > 255, i.e. n >= 23."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "source": [
+    "def machine_spin(x, fuel):\n",
+    "    for _ in range(fuel):\n",
+    "        if x == 0: return ('ret', None)\n",
+    "    return ('fuel', None)\n",
+    "def tr_spin(x): return ok(None) if x == 0 else DIV\n",
+    "agree = all(((machine_spin(x, F)[0] == 'ret') == (tr_spin(x).tag == 'ok')) and ((machine_spin(x, F)[0] == 'fuel') == (tr_spin(x).tag == 'div')) for x in U8 for F in (1, 10, 1000))\n",
+    "check(\"4a. spin: 'div' corresponds to fuel exhaustion for every tested fuel\", agree)\n",
+    "def machine_sum_to(n):\n",
+    "    s = i = 0\n",
+    "    while i < n:\n",
+    "        i += 1; s += i\n",
+    "        if s > 255: return ('panic', 'integerOverflow')\n",
+    "    return ('ret', s)\n",
+    "def tr_sum_loop(n, i=0, s=0):               # the recursive function Aeneas would emit\n",
+    "    if i < n:\n",
+    "        i2 = i + 1\n",
+    "        r = u8_add(s, i2)\n",
+    "        return bind(r, lambda s2: tr_sum_loop(n, i2, s2))\n",
+    "    return ok(s)\n",
+    "m = [machine_sum_to(n) for n in U8]\n",
+    "check(\"4b. sum_to: loop = recursive translation for all n in u8\", all(obs(tr_sum_loop(n)) == m[n] for n in U8))\n",
+    "first_bad = min(n for n in U8 if m[n][0] == 'panic')\n",
+    "check(\"4b. first panicking n is 23 (n(n+1)/2 = 276 > 255, 22*23/2 = 253 <= 255)\", first_bad == 23 and 22 * 23 // 2 == 253 and 23 * 24 // 2 == 276, f\"first_bad={first_bad}\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "[ok] 4a. spin: 'div' corresponds to fuel exhaustion for every tested fuel  \n",
+      "[ok] 4b. sum_to: loop = recursive translation for all n in u8  \n",
+      "[ok] 4b. first panicking n is 23 (n(n+1)/2 = 276 > 255, 22*23/2 = 253 <= 255)  first_bad=23\n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. A proof of `not P`: closed proposition versus negation under binders\n",
+    "Property of `avg`:  `P(a,b,r)`: `min(a,b) <= r <= max(a,b)`, evaluated on the translated function, with `Q(a,b) := exists r, avg a b = ok r /\\ P(a,b,r)`.\n",
+    "Closed statement  S := forall a b, H a b -> Q a b  (H trivially true here).   Its negation  not S  <=>  exists a b, H a b /\\ not Q a b.\n",
+    "Different, much stronger statement  S' := forall a b, H a b -> not Q a b.   Under binders, `not Q a b` for a free (a,b) is a claim about *every* instance."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "source": [
+    "def Q(a, b):\n",
+    "    r = tr_avg(a, b)\n",
+    "    return r.tag == 'ok' and min(a, b) <= r.v <= max(a, b)\n",
+    "H = lambda a, b: True\n",
+    "S      = all((not H(a, b)) or Q(a, b) for a, b in DOM)\n",
+    "notS   = any(H(a, b) and not Q(a, b) for a, b in DOM)\n",
+    "Sprime = all((not H(a, b)) or (not Q(a, b)) for a, b in DOM)\n",
+    "witnesses = [(a, b) for a, b in DOM if H(a, b) and not Q(a, b)]\n",
+    "check(\"5a. S is false and (not S) holds, certified by one witness\", (not S) and notS and witnesses[0] == (1, 255), f\"|W|={len(witnesses)}, first witness {witnesses[0]}\")\n",
+    "check(\"5a. the 'under binders' strengthening S' (forall x, not Q x) is false: Q holds at (0,0)\", (not Sprime) and Q(0, 0), f\"|{{Q}}|={len(DOM)-len(witnesses)}\")\n",
+    "check(\"5a. |W| = number of panicking inputs = 32640\", len(witnesses) == 32640)\n",
+    "# With a precondition H a b := a + b <= 255 the closed statement is TRUE and its negation is not provable\n",
+    "H2 = lambda a, b: a + b <= 255\n",
+    "check(\"5b. with H := a+b<=255, S_H holds (so no certificate for not S_H exists)\", all((not H2(a, b)) or Q(a, b) for a, b in DOM))\n",
+    "# Vacuity: if H is unsatisfiable, the under-binders form is provable and yet there is no witness at all\n",
+    "H3 = lambda a, b: a + b > 255 and a + b < 256\n",
+    "check(\"5c. unsatisfiable H: S'_H = forall x, H x -> not Q x holds vacuously, but the witness set is empty (no bug evidence)\",\n",
+    "      all((not H3(a, b)) or (not Q(a, b)) for a, b in DOM) and not any(H3(a, b) for a, b in DOM))\n",
+    "# And with H satisfiable and S'_H true (every H-input fails): S'_H implies the closed negation\n",
+    "H4 = lambda a, b: a + b >= 256\n",
+    "check(\"5c. satisfiable H := a+b>=256: S'_H holds (every such input fails) and so implies not S_H\", all((not H4(a, b)) or (not Q(a, b)) for a, b in DOM) and any(H4(a, b) for a, b in DOM))"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "[ok] 5a. S is false and (not S) holds, certified by one witness  |W|=32640, first witness (1, 255)\n",
+      "[ok] 5a. the 'under binders' strengthening S' (forall x, not Q x) is false: Q holds at (0,0)  |{Q}|=32896\n",
+      "[ok] 5a. |W| = number of panicking inputs = 32640  \n",
+      "[ok] 5b. with H := a+b<=255, S_H holds (so no certificate for not S_H exists)  \n",
+      "[ok] 5c. unsatisfiable H: S'_H = forall x, H x -> not Q x holds vacuously, but the witness set is empty (no bug evidence)  \n",
+      "[ok] 5c. satisfiable H := a+b>=256: S'_H holds (every such input fails) and so implies not S_H  \n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Reachability: the certificate is about `avg`, not about the program that calls it\n",
+    "Public caller: `fn mean_pct(a: u8, b: u8) -> u8 { avg(min(a, c), min(b, c)) }` for a clamp `c` (a property of the *caller*).\n",
+    "Reach_c = {(a', b'): a', b' <= c}.  Witnesses W from section 5 are real bugs of the program only if W intersect Reach_c is non-empty."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "source": [
+    "def count_formula(c):                       # pairs in [0,c]^2 with a+b >= 256\n",
+    "    return sum(max(0, min(c, 255) - max(0, 256 - a) + 1) for a in range(0, min(c, 255) + 1))\n",
+    "rows = []\n",
+    "for c in (50, 100, 127, 128, 150, 200, 255):\n",
+    "    brute = sum(1 for a, b in witnesses if a <= c and b <= c)\n",
+    "    rows.append((c, brute, count_formula(c)))\n",
+    "    assert brute == count_formula(c)\n",
+    "    print(f\"  clamp c={c:3d}: |W \u2229 Reach_c| = {brute:6d}   (closed form {count_formula(c)})\")\n",
+    "check(\"6. |W \u2229 Reach_c| equals the closed form for all tested c\", all(b == f for _, b, f in rows))\n",
+    "check(\"6. c <= 127: certificate holds for avg yet the program mean_pct never panics (false alarm)\", dict((c, b) for c, b, _ in rows)[127] == 0 and dict((c, b) for c, b, _ in rows)[128] == 1, \"c=127 -> 0, c=128 -> 1 (128+128=256)\")\n",
+    "# end-to-end: run the caller on its whole input domain\n",
+    "def mean_pct(a, b, c): return tr_avg(min(a, c), min(b, c))\n",
+    "for c in (100, 127, 128, 200):\n",
+    "    panics_c = sum(1 for a, b in DOM if mean_pct(a, b, c).tag == 'fail')\n",
+    "    print(f\"  mean_pct with clamp {c}: panicking inputs over all 65536 raw (a,b): {panics_c}\")\n",
+    "check(\"6. the caller with c=127 never fails on any of 65536 raw inputs; with c=128 it fails (a,b >= 128)\", all(mean_pct(a, b, 127).tag == 'ok' for a, b in DOM) and sum(1 for a, b in DOM if mean_pct(a, b, 128).tag == 'fail') == 128 * 128)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "  clamp c= 50: |W \u2229 Reach_c| =      0   (closed form 0)\n",
+      "  clamp c=100: |W \u2229 Reach_c| =      0   (closed form 0)\n",
+      "  clamp c=127: |W \u2229 Reach_c| =      0   (closed form 0)\n",
+      "  clamp c=128: |W \u2229 Reach_c| =      1   (closed form 1)\n",
+      "  clamp c=150: |W \u2229 Reach_c| =   1035   (closed form 1035)\n",
+      "  clamp c=200: |W \u2229 Reach_c| =  10585   (closed form 10585)\n",
+      "  clamp c=255: |W \u2229 Reach_c| =  32640   (closed form 32640)\n",
+      "[ok] 6. |W \u2229 Reach_c| equals the closed form for all tested c  \n",
+      "[ok] 6. c <= 127: certificate holds for avg yet the program mean_pct never panics (false alarm)  c=127 -> 0, c=128 -> 1 (128+128=256)\n",
+      "  mean_pct with clamp 100: panicking inputs over all 65536 raw (a,b): 0\n",
+      "  mean_pct with clamp 127: panicking inputs over all 65536 raw (a,b): 0\n",
+      "  mean_pct with clamp 128: panicking inputs over all 65536 raw (a,b): 16384\n",
+      "  mean_pct with clamp 200: panicking inputs over all 65536 raw (a,b): 29560\n",
+      "[ok] 6. the caller with c=127 never fails on any of 65536 raw inputs; with c=128 it fails (a,b >= 128)  \n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7. Translation soundness: wrong translations give wrong certificates in both directions\n",
+    "* U1 (too lax): unbounded integers.  Then `S` is *provable*, the real program violates it: a missed bug.\n",
+    "* U2 (wrapping, release-mode semantics): `Q` fails on the same inputs but as a wrong value, and panic-freedom becomes provable.  \"The bug\" depends on which Rust semantics is modelled.\n",
+    "* U3 (too strict): checks overflow on the *widened* sum `(a as u16 + b as u16)/2`, which is fine in Rust, but translates the u16 add as a u8 add: spurious `fail`, a false certificate."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "source": [
+    "def Q_with(trf): return lambda a, b: (lambda r: r.tag == 'ok' and min(a, b) <= r.v <= max(a, b))(trf(a, b))\n",
+    "trU1 = tr_naive_avg\n",
+    "trU2 = lambda a, b: ok(((a + b) % 256) // 2)\n",
+    "def rust_wide(a, b): return ('ret', (a + b) // 2)                  # `((a as u16 + b as u16) / 2) as u8` never panics\n",
+    "trU3 = lambda a, b: bind(u8_add(a, b), lambda s: ok(s // 2))      # wrong: widened add checked at u8\n",
+    "Wset = lambda q: sum(1 for a, b in DOM if not q(a, b))\n",
+    "print(\"  failing sets |{not Q}|: truth(avg, checked) =\", Wset(Q), \", U1 (unbounded) =\", Wset(Q_with(trU1)), \", U2 (wrapping) =\", Wset(Q_with(trU2)))\n",
+    "check(\"7a. U1: S is provable under the unbounded translation (no witness), but the program panics on 32640 inputs\", Wset(Q_with(trU1)) == 0 and Wset(Q) == 32640)\n",
+    "pf_checked = all(tr_avg(a, b).tag == 'ok' for a, b in DOM); pf_wrap = all(trU2(a, b).tag == 'ok' for a, b in DOM)\n",
+    "same_set = {(a, b) for a, b in DOM if not Q_with(trU2)(a, b)} == set(witnesses)\n",
+    "check(\"7b. U2: wrapping gives the SAME 32640 witnesses for P but for a different reason (wrong value, not panic); panic-freedom is provable under U2 and false under checked semantics\",\n",
+    "      same_set and pf_wrap and not pf_checked, f\"|W_wrap|={Wset(Q_with(trU2))}, panic-free: wrap={pf_wrap}, checked={pf_checked}\")\n",
+    "# wide version of the program: truth never panics, so S_wide holds; the strict translation certifies a bug\n",
+    "Qwide_true = lambda a, b: min(a, b) <= rust_wide(a, b)[1] <= max(a, b)\n",
+    "check(\"7c. U3: program `avg_wide` is correct on all inputs; the over-strict translation yields 32640 spurious witnesses\", all(Qwide_true(a, b) for a, b in DOM) and Wset(Q_with(trU3)) == 32640)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "  failing sets |{not Q}|: truth(avg, checked) = 32640 , U1 (unbounded) = 0 , U2 (wrapping) = 32640\n",
+      "[ok] 7a. U1: S is provable under the unbounded translation (no witness), but the program panics on 32640 inputs  \n",
+      "[ok] 7b. U2: wrapping gives the SAME 32640 witnesses for P but for a different reason (wrong value, not panic); panic-freedom is provable under U2 and false under checked semantics  |W_wrap|=32640, panic-free: wrap=True, checked=False\n",
+      "[ok] 7c. U3: program `avg_wide` is correct on all inputs; the over-strict translation yields 32640 spurious witnesses  \n"
+     ]
+    }
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Summary"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "source": [
+    "for n, ok_, d in CHECKS: print((\"PASS \" if ok_ else \"FAIL \") + n)\n",
+    "assert all(c[1] for c in CHECKS)\n",
+    "print(f\"\\n{len(CHECKS)} checks, all passed\")"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "PASS 1. monad laws on 260 values: left identity, right identity, associativity\n",
+      "PASS 2. naive translation disagrees with the machine on exactly the panicking inputs\n",
+      "PASS 2. Result translation agrees with the machine on all 65536 inputs\n",
+      "PASS 3a. bump: machine (store, outcome) = translation (returned value, Result) for all u8\n",
+      "PASS 3b. returned &mut: machine vs forward value + backward function, vec len<=3, i<=4, k<=2\n",
+      "PASS 4a. spin: 'div' corresponds to fuel exhaustion for every tested fuel\n",
+      "PASS 4b. sum_to: loop = recursive translation for all n in u8\n",
+      "PASS 4b. first panicking n is 23 (n(n+1)/2 = 276 > 255, 22*23/2 = 253 <= 255)\n",
+      "PASS 5a. S is false and (not S) holds, certified by one witness\n",
+      "PASS 5a. the 'under binders' strengthening S' (forall x, not Q x) is false: Q holds at (0,0)\n",
+      "PASS 5a. |W| = number of panicking inputs = 32640\n",
+      "PASS 5b. with H := a+b<=255, S_H holds (so no certificate for not S_H exists)\n",
+      "PASS 5c. unsatisfiable H: S'_H = forall x, H x -> not Q x holds vacuously, but the witness set is empty (no bug evidence)\n",
+      "PASS 5c. satisfiable H := a+b>=256: S'_H holds (every such input fails) and so implies not S_H\n",
+      "PASS 6. |W \u2229 Reach_c| equals the closed form for all tested c\n",
+      "PASS 6. c <= 127: certificate holds for avg yet the program mean_pct never panics (false alarm)\n",
+      "PASS 6. the caller with c=127 never fails on any of 65536 raw inputs; with c=128 it fails (a,b >= 128)\n",
+      "PASS 7a. U1: S is provable under the unbounded translation (no witness), but the program panics on 32640 inputs\n",
+      "PASS 7b. U2: wrapping gives the SAME 32640 witnesses for P but for a different reason (wrong value, not panic); panic-freedom is provable under U2 and false under checked semantics\n",
+      "PASS 7c. U3: program `avg_wide` is correct on all inputs; the over-strict translation yields 32640 spurious witnesses\n",
+      "\n",
+      "20 checks, all passed\n"
+     ]
+    }
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

code/notebooks/_lean_helper.py

@@ -0,0 +1,46 @@
+"""Tiny Lean runner for the 03xx notebooks. Goes ONLY through align.verify.lean.
+
+Rules (house + task): <= 2 workers (we use 1), nice 15, MemAvailable >= 45 GB checked
+before the pool starts, never the heavy-job lock, never touches a GPU.
+"""
+from __future__ import annotations
+import os, sys, re
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+sys.path.insert(0, "/apps/development/rlhf/align")
+from align.verify.lean import ReplPool, ReplLimits, find_toolchain
+from align.verify.lean.repl_pool import meminfo_available_gb
+
+HEADER = "import Mathlib\nimport Aesop\nset_option maxHeartbeats 400000\nopen Nat\n"
+
+class Lean:
+    def __init__(self, header: str = HEADER):
+        avail = meminfo_available_gb()
+        if avail < 45.0:
+            raise RuntimeError(f"MemAvailable {avail:.1f} GB < 45 GB: refusing to start Lean")
+        self.tc = find_toolchain()
+        self.pool = ReplPool.from_toolchain(self.tc, n_workers=1, limits=ReplLimits(nice=15))
+        self.header = header
+        self._cm = self.pool.acquire()
+        self.w = self._cm.__enter__()
+        self.env = self.w.header_env(header)
+
+    def run(self, code: str, timeout: float = 120.0, keep: bool = False):
+        """Run `code` on top of the header; return (messages, new_env)."""
+        resp = self.w.command(code, self.env, timeout)
+        if keep and "env" in resp:
+            self.env = resp["env"]
+        return resp
+
+    def show(self, code: str, timeout: float = 120.0, keep: bool = False) -> dict:
+        r = self.run(code, timeout, keep)
+        for m in r.get("messages", []):
+            print(f"[{m['severity']}] line {m['pos']['line']}: {m['data']}")
+        for s in r.get("sorries", []):
+            print("[sorries]", s.get("goal", "")[:200])
+        if not r.get("messages") and not r.get("sorries"):
+            print("[ok] no messages")
+        return r
+
+    def close(self):
+        try: self._cm.__exit__(None, None, None)
+        finally: self.pool.close()

code/notebooks/_m1_lib.py

@@ -0,0 +1,145 @@
+"""Shared helpers for the M1 notebooks (notes 0101-0110).
+
+Everything is fp64 CPU torch.  The "language model" is a deliberately tiny
+autoregressive policy whose *entire* sequence space can be enumerated, so every
+expectation in the notes can be computed EXACTLY (no Monte Carlo) and compared
+against the hand-derived formula to ~1e-15.  Monte Carlo checks are kept for the
+few identities that are about sampling (and are reported against 4 standard errors).
+
+Policy:  pi_theta(a | prev, t) = softmax(theta_tok[prev] + theta_pos[t])_a
+  - vocabulary {0,..,V-1} where token V-1 is EOS
+  - `prev` is the previous token (BOS = index V for t=0)
+  - a sequence stops at EOS or after L tokens (truncation, no EOS)
+This is a valid autoregressive factorisation  pi(o) = prod_t pi(o_t | o_<t)
+(with a Markov-1 context) so every identity about the general case holds verbatim.
+"""
+import os, sys, math, itertools
+os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")      # never touch the GPUs
+import torch
+torch.set_default_dtype(torch.float64)
+torch.set_num_threads(2)
+
+CHECKS = []
+
+def check(name, err, tol, kind="exact"):
+    """Record + assert a numerical identity.  kind='exact' (fp64 arithmetic) or 'mc' (4-sigma Monte Carlo)."""
+    err = float(err)
+    ok = err <= tol
+    CHECKS.append((name, err, tol, kind, ok))
+    print(f"CHECK|{kind}|{name}|{err:.3e}|tol={tol:.1e}|{'PASS' if ok else 'FAIL'}")
+    assert ok, f"{name}: err {err:.3e} > tol {tol:.1e}"
+
+def maxabs(a, b):
+    return float((torch.as_tensor(a) - torch.as_tensor(b)).abs().max())
+
+class TinyLM:
+    def __init__(self, V=4, L=4, seed=0, scale=1.0):
+        self.V, self.L = V, L
+        self.EOS = V - 1
+        g = torch.Generator().manual_seed(seed)
+        self.n_tok = (V + 1) * V          # rows: prev in {0..V-1, BOS}
+        self.n_pos = L * V
+        self.theta0 = scale * torch.randn(self.n_tok + self.n_pos, generator=g)
+        self.seqs, self.tokens, self.mask = self._enumerate()
+        self.N = len(self.seqs)
+        self.lengths = self.mask.sum(1)
+
+    # ---- enumeration of the whole sequence space --------------------------------
+    def _enumerate(self):
+        V, L, EOS = self.V, self.L, self.EOS
+        seqs = []
+        def rec(prefix):
+            if len(prefix) == L or (prefix and prefix[-1] == EOS):
+                seqs.append(tuple(prefix)); return
+            for a in range(V):
+                rec(prefix + [a])
+        rec([])
+        tokens = torch.full((len(seqs), L), EOS, dtype=torch.long)
+        mask = torch.zeros(len(seqs), L)
+        for i, s in enumerate(seqs):
+            tokens[i, :len(s)] = torch.tensor(s)
+            mask[i, :len(s)] = 1.0
+        return seqs, tokens, mask
+
+    def _logits_tables(self, theta):
+        V, L = self.V, self.L
+        tok = theta[:self.n_tok].reshape(V + 1, V)
+        pos = theta[self.n_tok:].reshape(L, V)
+        return tok, pos
+
+    def token_logprobs(self, theta):
+        """[N, L] log pi(o_t | o_<t) for every enumerated sequence (0 where masked)."""
+        V, L = self.V, self.L
+        tok, pos = self._logits_tables(theta)
+        prev = torch.cat([torch.full((self.N, 1), V, dtype=torch.long), self.tokens[:, :-1]], 1)  # BOS = V
+        logits = tok[prev] + pos[None, :, :]                       # [N, L, V]
+        lp = torch.log_softmax(logits, -1)
+        lp = lp.gather(-1, self.tokens[..., None]).squeeze(-1)
+        return lp * self.mask
+
+    def seq_logprob(self, theta):
+        return self.token_logprobs(theta).sum(1)
+
+    def probs(self, theta):
+        return self.seq_logprob(theta).exp()
+
+    # ---- naive, loop-based reference (used to cross-check the vectorised code) --
+    def seq_logprob_naive(self, theta, seq):
+        V = self.V
+        tok, pos = self._logits_tables(theta)
+        lp, prev = 0.0, V
+        for t, a in enumerate(seq):
+            logits = tok[prev] + pos[t]
+            z = logits - logits.max()
+            lp = lp + (z[a] - torch.log(torch.exp(z).sum()))
+            prev = a
+        return lp
+
+    # ---- exact state machinery for value functions ------------------------------
+    def prefixes(self):
+        """All non-terminal prefixes (the MDP states), shortest first."""
+        out = [()]
+        for k in range(1, self.L):
+            out += [p for p in itertools.product(range(self.V - 1), repeat=k)]  # non-EOS tokens only
+        return out
+
+    def cond_probs(self, theta, prefix):
+        V = self.V
+        tok, pos = self._logits_tables(theta)
+        prev = V if len(prefix) == 0 else prefix[-1]
+        return torch.softmax(tok[prev] + pos[len(prefix)], -1)
+
+def default_reward(lm):
+    """Binary verifier-style reward: success iff the sequence ends in EOS and its token sum is odd and length>=2."""
+    R = torch.zeros(lm.N)
+    for i, s in enumerate(lm.seqs):
+        if s[-1] == lm.EOS and len(s) >= 2 and (sum(s[:-1]) % 2 == 1):
+            R[i] = 1.0
+    return R
+
+def flat_grad(scalar, theta):
+    g, = torch.autograd.grad(scalar, theta, retain_graph=True, allow_unused=False)
+    return g
+
+def score_matrix(lm, theta):
+    """S[o] = grad_theta log pi(o)  as an [N, P] matrix (autograd, one backward per sequence)."""
+    th = theta.clone().requires_grad_(True)
+    lp = lm.seq_logprob(th)
+    rows = [torch.autograd.grad(lp[i], th, retain_graph=True)[0] for i in range(lm.N)]
+    return torch.stack(rows)
+
+def token_score_tensor(lm, theta):
+    """T[o, t] = grad_theta log pi(o_t | o_<t) as [N, L, P]."""
+    th = theta.clone().requires_grad_(True)
+    lp = lm.token_logprobs(th)
+    P = theta.numel()
+    out = torch.zeros(lm.N, lm.L, P)
+    for i in range(lm.N):
+        for t in range(int(lm.lengths[i])):
+            out[i, t] = torch.autograd.grad(lp[i, t], th, retain_graph=True)[0]
+    return out
+
+def perturb(theta, delta, seed=1):
+    g = torch.Generator().manual_seed(seed)
+    d = torch.randn(theta.numel(), generator=g)
+    return theta + delta * d / d.norm()

code/notebooks/_m1_run.py

@@ -0,0 +1,59 @@
+"""Execute a percent-format source (# %% / # %% [markdown]) and write an executed .ipynb with outputs.
+
+usage: python _m1_run.py src.py out.ipynb
+Runs all code cells in ONE namespace (like a kernel), captures stdout/stderr per cell,
+stores them as stream outputs.  Needs nothing but the stdlib (+ whatever the notebook imports).
+"""
+import sys, io, json, traceback, contextlib, os, time
+os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
+src, out = sys.argv[1], sys.argv[2]
+sys.path.insert(0, os.path.dirname(os.path.abspath(src)) + "/..")
+lines = open(src).read().split("\n")
+cells, cur, kind = [], [], None
+def flush():
+    global cur, kind
+    if kind is not None:
+        body = cur[:]
+        while body and body[-1] == "": body.pop()
+        cells.append((kind, body))
+    cur, kind = [], None
+for ln in lines:
+    if ln.startswith("# %%"):
+        flush(); kind = "markdown" if "[markdown]" in ln else "code"
+    else:
+        cur.append(ln)
+flush()
+ns = {"__name__": "__main__"}
+nb_cells, count = [], 0
+t0 = time.time()
+for kind, body in cells:
+    if kind == "markdown":
+        text = "\n".join(l[2:] if l.startswith("# ") else l.lstrip("#") for l in body)
+        nb_cells.append({"cell_type": "markdown", "metadata": {}, "source": text.splitlines(True)})
+        continue
+    count += 1
+    code = "\n".join(body)
+    so, se = io.StringIO(), io.StringIO()
+    err = None
+    with contextlib.redirect_stdout(so), contextlib.redirect_stderr(se):
+        try:
+            exec(compile(code, f"<cell {count}>", "exec"), ns)
+        except BaseException:
+            err = traceback.format_exc()
+    outputs = []
+    if so.getvalue(): outputs.append({"output_type": "stream", "name": "stdout", "text": so.getvalue().splitlines(True)})
+    if se.getvalue(): outputs.append({"output_type": "stream", "name": "stderr", "text": se.getvalue().splitlines(True)})
+    if err:
+        outputs.append({"output_type": "error", "ename": "Error", "evalue": err.splitlines()[-1], "traceback": err.splitlines()})
+    nb_cells.append({"cell_type": "code", "execution_count": count, "metadata": {}, "source": code.splitlines(True), "outputs": outputs})
+    if err:
+        print(err, file=sys.stderr)
+        nb = {"cells": nb_cells, "metadata": {}, "nbformat": 4, "nbformat_minor": 4}
+        json.dump(nb, open(out + ".FAILED", "w"), indent=1)
+        sys.exit(1)
+nb = {"cells": nb_cells,
+      "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
+                   "language_info": {"name": "python", "version": sys.version.split()[0]}},
+      "nbformat": 4, "nbformat_minor": 4}
+json.dump(nb, open(out, "w"), indent=1)
+print(f"wrote {out}: {count} code cells in {time.time()-t0:.1f}s")

code/notebooks/_m1_show.py

@@ -0,0 +1,4 @@
+import json,sys
+nb=json.load(open(sys.argv[1]))
+for c in nb['cells']:
+    for o in c.get('outputs',[]): print(''.join(o.get('text',o.get('traceback',''))),end='')

code/notes/STYLE_ALIGNMENT.md

@@ -0,0 +1,86 @@
+# Style alignment: align/notes and align/notebooks vs Flash-Attention-101
+
+Reference: `/apps/development/rlhf/reference-repos/Flash-Attention-101` (notes/0000-0008, 009, 010; README.md; scratch-pad.ipynb).
+Provenance caveat: the reference PDFs are handwritten notes whose title page (0000, p.1) is signed "Flash Attention from first principles, Umar Jamil". The repo (README clone URL viai957/Flash-Attention-101) packages them with the user's own README, notebooks and kernels. The "house style" below is therefore the handwritten-lecture style plus the repo's layout, not a typeset template. Where typeset equivalents are needed I mapped them (see "Mapping").
+
+## 1. Reference style rules (page citations are PDF page numbers)
+
+Layout and naming
+- R1. Numbered, one-topic-per-file notes named `NNNN - Title Case.pdf` (0001 - Multi Head Attention, 0003 - Online Softmax, ...). 0000 is a Table of Contents (0000 p.2 "Topics" list, p.3 "Prerequisites"). Numbering is not perfectly uniform (009, 010 have three digits).
+- R2. Length 2-9 pages per note (0002: 4, 0003: 6, 0004: 9, 0008: 9, 009/010: 2). One idea per page.
+- R3. Notebook: a single `scratch-pad.ipynb` that goes naive first (cells 0-9 "Conventional Softmax": loop for max, loop for normaliser, loop for division, then one consolidated typed function), then the improved version ("Safe Softmax" cell 10 on: single fused loop), then block/grid scaffolding for the Triton kernel. Markdown cell = pseudocode, code cell = direct transcription, then a consolidated function with a docstring. No asserts; checks are by printing. Kernel code lives in `triton/` and `cuda/`; README maps folder -> role.
+- R4. README: Overview, Folder Structure, "Understanding ...", "Implementation Details" in the order PyTorch naive -> Triton -> CUDA, Benchmarking, Setup.
+
+Page structure and tone
+- R5. Question-first opening. Each note starts with a bare problem statement as a question, then the pseudocode that fails: 0003 p.1 "The problem statement is: can we find a way to compute the softmax of a vector without going through the vector 3 times, but also preventing the exponential from exploding?"; 0004 p.2 "WHY SHOULD WE CARE?"; 0002 p.1 "But there's a problem! If the values of the vector are large, the exponential will explode!"; 0008 p.1 states the object ("Derivative: scalar input, scalar output") first.
+- R6. Motivation -> failing naive version -> fix -> proof. 0003: problem (p.1) -> try fusing (p.2) -> step-by-step trace on the same tiny vector x = [3, 2, 5, 1] (pp.2-4) -> "Can we prove this algorithm is correct? Let's do it!" (p.5) -> induction with base case and step, each closed by a marker (p.6).
+- R7. One running toy example reused on every page (x = [3,2,5,1]; the 8x8 blocked Q,K,V of 0004), with explicit shapes written next to every object ((N,d), (d,N), "ORIGINAL (8,128) / BLOCK (4,128)", 0002 p.1, 0004 pp.2-5).
+- R8. Pseudocode appears next to the derivation, is restated in full each time it is modified, and is labelled "Pseudocode" / "New pseudocode" (0003 pp.1, 4).
+- R9. Figures carry the argument: grid pictures of blocked matrices, computation graph for autograd (0008 pp.6-7), memory-hierarchy figure (0000 p.1). Colour has a fixed role: blue = given/derivation, red = warning or the open question, purple = annotation/shape, green = "correct" / the question to prove.
+- R10. Warnings are explicit and loud: "WRONG!" box with the reason (0004 p.5), "The l3 we computed is wrong! Here it is wrong!" (0003 p.3). A "THE IDEA" boxed summary precedes the fix (0004 p.6).
+- R11. Tone: second person plural, conversational ("Let's restore it... with a twist!", "Let's multiply by V", "Lots of patience", 0000 p.3). Short sentences, no citations inside notes; the only external reference is a pasted algorithm (FlashAttention-2 Algorithm 1, 0004 p.9) annotated by hand.
+- R12. Complexity is stated per step (time and memory reads, 0002 p.3).
+- R13. Notation: m_i (running max), l_i (normaliser), O_i, S = QK^T, P = softmax(S), blocks Q_i, K_j, V_j; softmax* for the unnormalised block softmax (0004 p.4); upstream/local/downstream gradient vocabulary (0008 p.9).
+
+## 2. Mapping to the typeset notes (what "same style" means here)
+
+| Reference | Typeset equivalent used |
+|---|---|
+| R1 file naming | already `NNNN - Title.pdf`; kept |
+| R5 question first | `problem` box titled "The problem" whose first line is `\question{...}` |
+| R6 motivation -> math -> proof | Notation, Derivation (proofs inline), unchanged content |
+| R7 shapes and a running toy | the notes already carry toy numbers in `\ver{...}` lines and notation tables; not altered |
+| R10 warnings | `What breaks` box (all 24) |
+| R9 "THE IDEA" boxed summary | `Boxed result` box (all 24) |
+| (flow to next page) | `Next` box (all 24) |
+
+## 3. Gaps found (before the change)
+
+Compared 0107 (M1), 0203 (M2), 0304 (M3) and then grepped all 24.
+
+Three series had three different style files and three different skeletons:
+- Header: M1 `\hdr` (no "Note" word, grey series line), M2/M3 `\notehead` ("Note NNNN -- ..."). 
+- Opening: M1 grey "The problem this solves" box; M2 numbered section "The problem this solves"; M3 numbered section "The problem this concept solves". None opened with an explicit question (R5), although many posed one inside the paragraph.
+- Result: M1 yellow "Result"; M2 blue box; M3 blue box, section called "Boxed final results" (M3) vs "Boxed result" (M1 box with no section, M2).
+- Failure modes: M1 box "What breaks" under a section of the same name; M2 box with title override; M3 plain section (0301 titled "What breaks (failure modes and the assumption each violates)").
+- Next: M1 blue box; M2/M3 bold inline line under "Link to the next note" (M2) or none (M3).
+- References: M2 had `\section{References}` followed by `thebibliography` (duplicate heading).
+- Page: M1 1in margins, 11pt, colour links; M2 2.2cm; M3 2.2cm with Unicode/Lean listings; M2 10pt, M1/M3 11pt.
+- Macro clashes between the files: `\ind` (bracketed one-argument in M1, bare symbol in M2), `\Pr` (operator vs blackboard P), `\nextnote` (environment vs command).
+
+Not present in any of the 24 (and not added, because it would be new content): hand-drawn style figures (R9), a "WRONG!" worked-trace page (R10), complexity-per-step tables (R12), a Table of Contents note (R1 0000), a prerequisites page.
+
+## 4. What I changed
+
+Style infrastructure
+- `notes/src/common.sty` replaces `common-m1.sty`, `common-m2.sty`, `common-m3.sty` (the three are deleted; originals backed up in `/home/vignesh/.claude/jobs/b952890a/tmp/notes_style/backup/`). It is the union of the three: geometry 2.2cm, 11pt documents, `hidelinks`, Cerulean result box, BrickRed breaks box, grey problem box, blue Next box, Lean `listings` and the Unicode declarations, all shared macros.
+- `\ind` is the bare symbol everywhere; the bracketed M1 form is now `\indb{...}` (renamed in 0101-0110 only; same rendering as before). `\Pr` is `\mathbb{P}` in all notes (0301, 0302, 0304, 0305 used the default operator "Pr"; only the glyph changes).
+- `notes/src/build_all.sh` rebuilds all 24 (or a subset: `./build_all.sh 0107 0203`), takes the PDF file name from the existing `../NNNN - *.pdf`. `build.sh` and `build_m1.sh` still work unchanged.
+
+All 24 `.tex` files
+- `\usepackage{common}`, `\documentclass[11pt]`.
+- One header, `\notehead{NNNN}{title}{subtitle}` (M1's `\hdr` converted; subtitle text kept: series line plus companion notebook path).
+- One opening: a grey "The problem" box whose first line is a new one-sentence `\question{...}` (24 questions, each only restating the question the existing paragraph already poses; no new claim, no number), followed by the original paragraph untouched. For M2/M3 the old numbered section was converted into this box.
+- One tail: `Boxed result` section, `What breaks` section with the red box (M2/M3 explicit title override dropped, the default is the same text), `Next` box, `References` (M2/M3 duplicate heading removed).
+- Section numbers: the box does not take a number in M1 (Notation = 1, Derivation = 2). In M2/M3 the old problem section was section 1, so I inserted `\stepcounter{section}` after the box so Notation stays 2 and Derivation 3. Reason: many text and notebook cross-references use those numbers (for example 0301 "\S3.3", 0302 notebook "section 3.4", 0304 notebook "section 3", 0203 notebook "section 3", 0102 cited as "Sec.~2.5" from 0105). Renumbering needs a coordinated edit of notes and notebooks and was not done. Consequence: M1 numbers sections from 1, M2/M3 from 2.
+
+Rebuild
+- 24 of 24 PDFs rebuilt with latexmk, 0 failures, 0 undefined references/citations. Page counts: 0103 5->4, 0201 5->6, 0202 5->6, 0204 5->6, 0205 5->6, 0206 5->6, 0207 5->6, 0302 6->7, 0304 6->7, all others unchanged. Text diff of old vs new PDFs (word multiset) shows only header, section-number, heading, box-title and the 24 added `The question.` lines; no equation, number or verified-result line changed.
+- `notes/src/build/m1/` (stale build artefacts of the old M1 script) was deleted; `notes/src/build/all/` is the new build dir.
+
+Notebooks
+- Not changed. They already follow the reference's naive -> efficient progression with asserts and verification tables (0107 section 7 "Naive vs efficient implementation"; 0203 sections 1-3 naive exact/float vs stable form; 0304 "A1 NAIVE ... A2 EFFICIENT"). The reference notebook has no asserts and a single `scratch-pad.ipynb`; ours are per-note numbered `NNNN-slug.ipynb`, which I kept (the numbered naming matches R1 and the .tex/PDF linkage). Nothing was re-executed; no Lean notebook touched.
+
+## 5. Not aligned (left deliberately)
+
+1. Handwritten diagram/colour language (R9, R10), worked trace on one toy vector (R7), per-step complexity (R12): would be new content. A reasonable follow-up is a `\begin{warning}` "WRONG!" box and a 4-element trace per note, authored by hand per topic.
+2. 0000 Table of Contents note and a prerequisites page: not requested; a generated `0000 - Table of Contents.pdf` from the 24 titles is mechanical if wanted.
+3. Section numbering offset between M1 and M2/M3 (see above).
+4. M3 notes keep their `\paragraph{Numerical verification (notebook ...)}` block while M1 uses inline `\ver{...}` tags and M2 mixes both; both carry the same information, and merging them touches content layout.
+5. Series subtitle wording still differs ("Batch M1/M2", M3 has descriptive text).
+6. Notebook per-note files vs reference single scratch-pad: kept.
+
+## 6. Possible content issues (not edited)
+
+- `notebooks/0101-autoregressive-likelihood-sft.ipynb` cell 2 cites "(note, Sec. 3.1)" for the proper-distribution lemma; in the M1 numbering the lemma sits under Section 2 (Derivation, "Step 2 (Lemma 1)"), so this reference looks stale. Not verified against the notebook text beyond the grep.
+- No mathematical claim or number was found to be wrong during this pass (this was a style pass; the content was not re-derived).

code/scripts/bench_7b_end_to_end.py

@@ -0,0 +1,283 @@
+"""Benchmark #2: end-to-end step time on the REAL 7B, triton_fa2 vs SDPA.
+
+Run (GPU 3 window; pin BEFORE torch import):
+    CUDA_VISIBLE_DEVICES=3 TMPDIR=/apps/development/rlhf/nano-mistral-o1/.tmp \
+        uv run python scripts/bench_7b_end_to_end.py
+
+WHAT THIS MEASURES, AND WHAT IT DELIBERATELY DOES NOT
+-----------------------------------------------------
+A full-parameter 7B training step does NOT fit on one A10G: bf16 weights are
+13.5 GiB and bf16 parameter gradients are another 13.5 GiB, so 27 GiB against a
+23 GiB card. Real SFT gets around this with 4-way FSDP plus an offloaded
+optimizer (section-07/11); that is not reproducible inside a single-GPU window,
+and pretending otherwise would be measuring a different thing than the label.
+
+So this benchmark freezes every parameter and requires grad on `inputs_embeds`
+instead. Backward still flows through every attention and MLP in all 32 layers,
+computing the full activation-gradient chain -- dQ/dK/dV included, which is
+exactly the part the attention backend controls. What is skipped is the
+per-parameter weight-gradient GEMM, which is IDENTICAL for both backends.
+
+The consequence is worth stating precisely, because it decides how to read the
+result: the absolute per-step delta (triton minus sdpa, in ms) is exactly what
+it would be in a full step, since the omitted work is backend-independent and
+additive. The delta as a PERCENTAGE is inflated, because the denominator is
+missing that same backend-independent work. **This benchmark therefore reports
+an UPPER BOUND on the true percentage cost of choosing triton_fa2.** A number
+that bounds the answer from the safe side is worth more here than a number that
+requires four GPUs and a training run to obtain.
+
+Both no-checkpointing and activation-checkpointing configurations are measured:
+under checkpointing the attention forward runs a second time during backward,
+so the backend-attributable share is larger, and SFT is expected to checkpoint.
+"""
+
+from __future__ import annotations
+
+import glob
+import json
+import os
+import sys
+import time
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from align.utils.guard import apply_resource_guards
+
+# Light job by the box rules: no heavy lock (other GPU jobs are live), no
+# dataloaders, no multi-GB writes.
+apply_resource_guards(cpu_threads=8, heavy=False, label="bench-7b-attn")
+
+import torch
+from torch.utils.checkpoint import checkpoint as ckpt_fn
+
+from align.model import checkpoint as ckpt
+from align.model.mistral import NmoMistralForCausalLM
+
+REPO = "/apps/development/rlhf/nano-mistral-o1"
+OUT = f"{REPO}/docs/bench_7b_end_to_end.json"
+SNAPSHOT = (
+    "/apps/voxtral/.cache/huggingface/hub/"
+    "models--mistralai--Mistral-7B-v0.3/snapshots/"
+    "caa1feb0e54d415e2df31207e5f4e273e33509b1"
+)
+
+SEQ = 1024
+BATCH = 1
+#: Hard wall-clock budget. The GPU window is shared and SFT launches the moment
+#: the golden gate passes; a benchmark that overstays blocks a training run.
+DEADLINE_S = float(os.environ.get("NMO_BENCH_DEADLINE_S", 1500))
+_START = time.time()
+
+
+def _left() -> float:
+    return DEADLINE_S - (time.time() - _START)
+
+
+def _load_7b(attn_impl: str, state_dict, cfg):
+    model = NmoMistralForCausalLM.from_state_dict(cfg, state_dict, attn_implementation=attn_impl)
+    model = model.to(torch.bfloat16).to("cuda").eval()
+    # Freeze: no 13.5 GiB parameter-gradient buffer. See module docstring.
+    for p in model.parameters():
+        p.requires_grad_(False)
+    return model
+
+
+def _blocks(model):
+    inner = model.model if hasattr(model, "model") else model
+    for name in ("layers", "blocks", "h"):
+        if hasattr(inner, name):
+            return getattr(inner, name)
+    raise AttributeError("cannot find the decoder block list on the model")
+
+
+def _timed(fn, iters, warmup):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(iters):
+        fn()
+    torch.cuda.synchronize()
+    return (time.perf_counter() - t0) / iters * 1e3
+
+
+def _run_config(model, embeds, labels, use_ckpt: bool, iters: int):
+    """fwd-only and fwd+bwd timings for one model/checkpointing combination."""
+    blocks = _blocks(model)
+    originals = [b.forward for b in blocks]
+    if use_ckpt:
+        for b in blocks:
+            fwd = b.forward
+            b.forward = lambda *a, _f=fwd, **kw: ckpt_fn(_f, *a, use_reentrant=False, **kw)
+
+    try:
+
+        def fwd_only():
+            with torch.no_grad():
+                model(inputs_embeds=embeds.detach())
+
+        def fwd_bwd():
+            if embeds.grad is not None:
+                embeds.grad = None
+            logits = model(inputs_embeds=embeds)
+            loss = torch.nn.functional.cross_entropy(
+                logits.float().view(-1, logits.shape[-1]), labels.view(-1)
+            )
+            loss.backward()
+
+        torch.cuda.reset_peak_memory_stats()
+        f = _timed(fwd_only, iters, 2)
+        fb = _timed(fwd_bwd, iters, 2)
+        peak = torch.cuda.max_memory_allocated() / 2**30
+    finally:
+        # `b.forward = original_bound_method` would leave the instance holding a
+        # bound method that holds the instance -- a reference cycle the caching
+        # allocator cannot reclaim on `empty_cache()`, which OOM'd the second
+        # model load. Deleting the instance attribute restores the class method
+        # and leaves no cycle.
+        for b in blocks:
+            if "forward" in b.__dict__:
+                del b.__dict__["forward"]
+        del originals
+    return {"fwd_ms": f, "fwd_bwd_ms": fb, "peak_gib": peak}
+
+
+def main() -> int:
+    if torch.cuda.device_count() != 1:
+        print(
+            f"{torch.cuda.device_count()} devices visible; pin exactly one "
+            "(shared box, three other GPU jobs live)",
+            file=sys.stderr,
+        )
+        return 1
+    free_gib = torch.cuda.mem_get_info()[0] / 2**30
+    print(f"device: {torch.cuda.get_device_name(0)}  free {free_gib:.1f} GiB")
+    if free_gib < 20.0:
+        print(
+            "refusing: need ~20 GiB free for 13.5 GiB of bf16 weights plus "
+            "activations; another job is resident on this device.",
+            file=sys.stderr,
+        )
+        return 1
+
+    from align.eval.golden import resolve_model_config
+
+    cfg = resolve_model_config(SNAPSHOT)
+    print(
+        f"config: {cfg.n_layers}L hidden={cfg.hidden_size} "
+        f"{cfg.n_heads}q:{cfg.n_kv_heads}kv head_dim={cfg.head_dim} "
+        f"vocab={cfg.vocab_size} sliding_window={cfg.sliding_window}"
+    )
+
+    import safetensors.torch as st
+
+    t0 = time.time()
+    hf_sd = {}
+    for shard in sorted(glob.glob(os.path.join(SNAPSHOT, "*.safetensors"))):
+        hf_sd.update(st.load_file(shard))
+    nmo_sd = ckpt.hf_to_nmo(hf_sd, expected_n_layers=cfg.n_layers)
+    del hf_sd
+    print(f"loaded + remapped {len(nmo_sd)} tensors in {time.time() - t0:.1f}s")
+
+    torch.manual_seed(0)
+    labels = torch.randint(0, cfg.vocab_size, (BATCH, SEQ), device="cuda")
+    results = {}
+
+    # One backend per PROCESS when an arg is given. 13.5 GiB of weights leaves
+    # no room for a second model, and process exit is the only VRAM release
+    # that cannot be defeated by a stray reference. Partials are merged below.
+    only = sys.argv[1] if len(sys.argv) > 1 else None
+    impls = (only,) if only else ("sdpa", "triton_fa2")
+
+    for impl in impls:
+        if _left() < 240:
+            print(f"\ndeadline guard: {_left():.0f}s left, skipping {impl}", file=sys.stderr)
+            break
+        print(f"\n=== {impl} ===", flush=True)
+        model = _load_7b(impl, nmo_sd, cfg)
+        embeds = torch.randn(
+            BATCH, SEQ, cfg.hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True
+        )
+        results[impl] = {}
+        for use_ckpt in (False, True):
+            if _left() < 120:
+                print(
+                    f"deadline guard: {_left():.0f}s left, skipping {impl}/ckpt={use_ckpt}",
+                    file=sys.stderr,
+                )
+                break
+            r = _run_config(model, embeds, labels, use_ckpt, iters=5)
+            key = "activation_ckpt" if use_ckpt else "no_ckpt"
+            results[impl][key] = r
+            print(
+                f"  {key:16s} fwd {r['fwd_ms']:8.1f} ms   "
+                f"fwd+bwd {r['fwd_bwd_ms']:8.1f} ms   peak {r['peak_gib']:5.2f} GiB",
+                flush=True,
+            )
+        del model, embeds
+        import gc
+
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    # Merge any partial written by a sibling process, so running one backend
+    # per process still produces a single complete receipt.
+    partial = f"{REPO}/.tmp/bench7b_partial.json"
+    if only:
+        prev = {}
+        if os.path.isfile(partial):
+            with open(partial) as f:
+                prev = json.load(f)
+        prev.update(results)
+        with open(partial, "w") as f:
+            json.dump(prev, f, indent=2)
+        results = prev
+        print(f"\npartial results for {sorted(results)} -> {partial}")
+
+    payload = {
+        "note": "parameters FROZEN; grad flows via inputs_embeds. Absolute ms "
+        "deltas are exact; percentages are an UPPER BOUND on the true "
+        "cost, because the omitted weight-grad GEMMs are "
+        "backend-independent and would enlarge the denominator only. "
+        "See the module docstring.",
+        "device": torch.cuda.get_device_name(0),
+        "torch": torch.__version__,
+        "model": "mistralai/Mistral-7B-v0.3 (real weights)",
+        "batch": BATCH,
+        "seq": SEQ,
+        "dtype": "bfloat16",
+        "results": results,
+    }
+
+    if "sdpa" in results and "triton_fa2" in results:
+        payload["comparison"] = {}
+        print("\n=== triton_fa2 vs sdpa on the real 7B ===")
+        for key in ("no_ckpt", "activation_ckpt"):
+            if key in results["sdpa"] and key in results["triton_fa2"]:
+                s = results["sdpa"][key]["fwd_bwd_ms"]
+                t = results["triton_fa2"][key]["fwd_bwd_ms"]
+                c = {
+                    "sdpa_ms": s,
+                    "triton_ms": t,
+                    "delta_ms": t - s,
+                    "ratio": t / s,
+                    "upper_bound_pct": 100.0 * (t - s) / s,
+                }
+                payload["comparison"][key] = c
+                print(
+                    f"  {key:16s} sdpa {s:8.1f} ms -> triton {t:8.1f} ms "
+                    f"({c['delta_ms']:+.1f} ms, {c['upper_bound_pct']:+.1f}% "
+                    "UPPER BOUND)"
+                )
+
+    os.makedirs(os.path.dirname(OUT), exist_ok=True)
+    with open(OUT, "w") as f:
+        json.dump(payload, f, indent=2)
+        f.write("\n")
+    print(f"\nwrote {OUT}  ({_left():.0f}s of window left)")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

code/scripts/bench_flash_attention.py

@@ -0,0 +1,234 @@
+"""Kernel-level benchmark: triton_fa2 vs torch SDPA at the real Mistral shapes.
+
+Run (pin ONE free device; this is a light job and must NOT take the heavy lock):
+    CUDA_VISIBLE_DEVICES=<free gpu> uv run python scripts/bench_flash_attention.py
+
+Shapes are Mistral-7B-v0.3's: b=4, 32 query heads : 8 KV heads, head_dim 128,
+causal, bf16, seq {512, 1024, 2048}. The forward column is compared against the
+baseline pinned in the plan artifact `tmp/sdpa_baseline.json`, measured on this
+same box: 0.189 / 0.624 / 2.282 ms.
+
+Purpose is to bound a decision, not to win: research (FA2 paper 2307.08691,
+flash-attention#1556) says Triton loses to CUTLASS at head_dim 128 on Ampere,
+~1.3-1.5x forward and ~2x backward, and head_dim 128 is exactly Mistral's. At
+seq 1024 attention is a few percent of layer FLOPs, so even a 2x backward loss
+is a small single-digit share of step time. This script turns that into a
+measured number so the SDPA-default decision is evidence rather than vibes.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from align.utils.guard import apply_resource_guards
+
+# Light job: no heavy lock, no 7B model. Still honors the path policy so no
+# cache lands off /apps.
+apply_resource_guards(cpu_threads=8, heavy=False, label="bench-flash-attention")
+
+import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+from align.kernels.flash_attention import flash_attention
+
+REPO = "/apps/development/rlhf/align"
+OUT = f"{REPO}/docs/bench_flash_attention.json"
+
+#: Measured on this box during planning; see the module docstring.
+PINNED_SDPA_FWD_MS = {512: 0.189, 1024: 0.624, 2048: 2.282}
+
+B, HQ, HKV, D = 4, 32, 8, 128
+SEQS = (512, 1024, 2048)
+
+
+def _bench(fn, iters=30, warmup=10):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iters):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) / iters
+
+
+def bench_layer_share(seq: int = 1024, batch: int = 2) -> dict:
+    """Put attention back in context: one full Mistral decoder layer, fwd+bwd.
+
+    Benchmark 1 gives the ratio ON the attention kernel; this converts it into
+    the number the SDPA-vs-triton decision actually turns on -- share of a
+    training step. A layer is built from Mistral-7B-v0.3's real dimensions
+    (hidden 4096, intermediate 14336, 32:8 heads, head_dim 128) so the MLP and
+    projection GEMMs are correctly sized against the attention core.
+
+    This is NOT the section-11/12 end-to-end number the section also asks for;
+    that one needs the real 7B weights and a GPU window, and is deferred to it.
+    A single layer times 32 is a tight bound on the attention share because the
+    per-layer work is identical across layers.
+    """
+    import torch.nn as nn
+
+    H, I = 4096, 14336
+    dev_t, dt = "cuda", torch.bfloat16
+    lin = lambda a, b: nn.Linear(a, b, bias=False, device=dev_t, dtype=dt)
+    wq, wk, wv, wo = lin(H, HQ * D), lin(H, HKV * D), lin(H, HKV * D), lin(HQ * D, H)
+    wg, wu, wd = lin(H, I), lin(H, I), lin(I, H)
+    x = torch.randn(batch * seq, H, device=dev_t, dtype=dt, requires_grad=True)
+
+    def layer(attn_fn):
+        def run():
+            x.grad = None
+            q = wq(x).view(batch, seq, HQ, D).transpose(1, 2)
+            k = wk(x).view(batch, seq, HKV, D).transpose(1, 2)
+            v = wv(x).view(batch, seq, HKV, D).transpose(1, 2)
+            a = attn_fn(q, k, v)
+            h = x + wo(a.transpose(1, 2).reshape(batch * seq, H))
+            out = h + wd(torch.nn.functional.silu(wg(h)) * wu(h))
+            out.sum().backward()
+
+        return run
+
+    def sdpa_attn(q, k, v):
+        with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
+            return torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, is_causal=True, enable_gqa=True
+            )
+
+    t_sdpa = _bench(layer(sdpa_attn), iters=20, warmup=8)
+    t_tri = _bench(layer(lambda q, k, v: flash_attention(q, k, v, causal=True)), iters=20, warmup=8)
+
+    out = {
+        "seq": seq,
+        "batch": batch,
+        "hidden": H,
+        "intermediate": I,
+        "n_layers": 32,
+        "layer_sdpa_ms": t_sdpa,
+        "layer_triton_ms": t_tri,
+        "layer_ratio": t_tri / t_sdpa,
+        "step_penalty_pct": 100.0 * (t_tri - t_sdpa) / t_sdpa,
+        "est_32layer_sdpa_s": t_sdpa * 32 / 1000.0,
+        "est_32layer_triton_s": t_tri * 32 / 1000.0,
+    }
+    print(f"\nlayer-level (Mistral decoder layer, b={batch} seq={seq}, fwd+bwd):")
+    print(f"  SDPA   {t_sdpa:7.2f} ms/layer -> {out['est_32layer_sdpa_s']:.2f} s / 32 layers")
+    print(f"  triton {t_tri:7.2f} ms/layer -> {out['est_32layer_triton_s']:.2f} s / 32 layers")
+    print(f"  cost of choosing triton_fa2: {out['step_penalty_pct']:+.1f}% of the 32-layer body")
+    return out
+
+
+def main() -> int:
+    if not torch.cuda.is_available():
+        print("no CUDA device visible", file=sys.stderr)
+        return 1
+    if torch.cuda.device_count() != 1:
+        print(
+            f"{torch.cuda.device_count()} devices visible; pin exactly one "
+            "with CUDA_VISIBLE_DEVICES (shared box)",
+            file=sys.stderr,
+        )
+        return 1
+
+    dev = torch.cuda.get_device_name(0)
+    cc = ".".join(str(x) for x in torch.cuda.get_device_capability(0))
+    print(f"device: {dev} (sm{cc.replace('.', '')})  torch {torch.__version__}")
+    print(f"shapes: b={B}, {HQ}q:{HKV}kv heads, head_dim={D}, causal, bf16\n")
+
+    results = {}
+    for seq in SEQS:
+        q = torch.randn(B, HQ, seq, D, device="cuda", dtype=torch.bfloat16)
+        k = torch.randn(B, HKV, seq, D, device="cuda", dtype=torch.bfloat16)
+        v = torch.randn(B, HKV, seq, D, device="cuda", dtype=torch.bfloat16)
+        qg, kg, vg = (t.clone().requires_grad_() for t in (q, k, v))
+        dout = torch.randn_like(qg)
+
+        def sdpa_fwd():
+            with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
+                return torch.nn.functional.scaled_dot_product_attention(
+                    q, k, v, is_causal=True, enable_gqa=True
+                )
+
+        def tri_fwd():
+            return flash_attention(q, k, v, causal=True)
+
+        def sdpa_fwd_bwd():
+            for t in (qg, kg, vg):
+                t.grad = None
+            with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
+                o = torch.nn.functional.scaled_dot_product_attention(
+                    qg, kg, vg, is_causal=True, enable_gqa=True
+                )
+            o.backward(dout)
+
+        def tri_fwd_bwd():
+            for t in (qg, kg, vg):
+                t.grad = None
+            flash_attention(qg, kg, vg, causal=True).backward(dout)
+
+        row = {
+            "sdpa_fwd_ms": _bench(sdpa_fwd),
+            "triton_fwd_ms": _bench(tri_fwd),
+            "sdpa_fwd_bwd_ms": _bench(sdpa_fwd_bwd),
+            "triton_fwd_bwd_ms": _bench(tri_fwd_bwd),
+            "pinned_sdpa_fwd_ms": PINNED_SDPA_FWD_MS[seq],
+        }
+        row["fwd_ratio_triton_over_sdpa"] = row["triton_fwd_ms"] / row["sdpa_fwd_ms"]
+        row["fwd_bwd_ratio_triton_over_sdpa"] = row["triton_fwd_bwd_ms"] / row["sdpa_fwd_bwd_ms"]
+        row["sdpa_fwd_drift_vs_pinned"] = row["sdpa_fwd_ms"] / row["pinned_sdpa_fwd_ms"]
+        results[str(seq)] = row
+        # Rebind to None rather than `del`: the bench closures above hold these
+        # names as free variables, and `del` empties the closure cells (ruff
+        # F821; a later call would NameError). Rebinding frees the tensors just
+        # the same, so empty_cache() can actually reclaim.
+        q = k = v = qg = kg = vg = dout = None
+        torch.cuda.empty_cache()
+
+    hdr = (
+        f"{'seq':>6} | {'SDPA fwd':>9} {'pinned':>8} | {'triton fwd':>11} "
+        f"{'ratio':>7} | {'SDPA f+b':>9} {'triton f+b':>11} {'ratio':>7}"
+    )
+    print(hdr)
+    print("-" * len(hdr))
+    for seq in SEQS:
+        r = results[str(seq)]
+        print(
+            f"{seq:>6} | {r['sdpa_fwd_ms']:>9.3f} {r['pinned_sdpa_fwd_ms']:>8.3f} | "
+            f"{r['triton_fwd_ms']:>11.3f} {r['fwd_ratio_triton_over_sdpa']:>6.2f}x | "
+            f"{r['sdpa_fwd_bwd_ms']:>9.3f} {r['triton_fwd_bwd_ms']:>11.3f} "
+            f"{r['fwd_bwd_ratio_triton_over_sdpa']:>6.2f}x"
+        )
+
+    layer = bench_layer_share()
+
+    payload = {
+        "device": dev,
+        "compute_capability": cc,
+        "torch": torch.__version__,
+        "shape": {
+            "batch": B,
+            "q_heads": HQ,
+            "kv_heads": HKV,
+            "head_dim": D,
+            "causal": True,
+            "dtype": "bfloat16",
+        },
+        "results": results,
+        "layer_level": layer,
+    }
+    os.makedirs(os.path.dirname(OUT), exist_ok=True)
+    with open(OUT, "w") as f:
+        json.dump(payload, f, indent=2)
+        f.write("\n")
+    print(f"\nwrote {OUT}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

code/scripts/build_data.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python
+"""Build every corpus this project trains on, in the only legal order.
+
+    python scripts/build_data.py --out <build root> [--tokenizer <dir>] [--curation-target N]
+
+Stages: eval sets -> decontamination hash set -> SFT -> Arm A -> Arm B.
+Each refuses to run without its upstream receipt; the final manifest
+(`<out>/manifest.json`) pins every content hash, every measured count against
+its pinned value, reject and contamination counts, the length-bias baseline,
+the delimiter ids, the tokenizer hash, and the dataset revisions. Sections
+11/12 quote that manifest in their launch banner and refuse to start without
+it.
+
+Host safety (this box, 186GB, NO swap, shared): `apply_resource_guards(
+heavy=True)` runs the preflight and takes the exclusive job flock BEFORE
+anything heavy, so this build cannot run concurrently with a training job --
+that exact overlap killed the predecessor's first SFT run. Every corpus is
+streamed; nothing is held resident.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from align.utils.guard import apply_resource_guards  # noqa: E402
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger("build_data")
+
+DEFAULT_OUT = "/apps/development/rlhf/nano-mistral-o1/data"
+TOKENIZER_SNAPSHOT = (
+    "/apps/voxtral/.cache/huggingface/hub/models--mistralai--Mistral-7B-v0.3/"
+    "snapshots/caa1feb0e54d415e2df31207e5f4e273e33509b1"
+)
+#: Byte-identical tokenizer files vendored in-repo (assets/tokenizer/.../PROVENANCE.md).
+#: The HF-cache snapshot is NOT on every box; the build must not die on a missing
+#: cache when the same four files are in the repo.
+VENDORED_TOKENIZER = str(
+    Path(__file__).resolve().parent.parent / "assets" / "tokenizer" / "mistral-7b-v0.3"
+)
+
+
+def default_tokenizer_dir() -> str:
+    """The HF snapshot if it holds a tokenizer.json, else the vendored copy."""
+    for c in (TOKENIZER_SNAPSHOT, VENDORED_TOKENIZER):
+        if (Path(c) / "tokenizer.json").is_file():
+            return c
+    return TOKENIZER_SNAPSHOT  # fails loudly in load_tokenizer, naming the path
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument(
+        "--out", default=DEFAULT_OUT, help="build root; every path under it is asserted onto /apps"
+    )
+    ap.add_argument("--tokenizer", default=default_tokenizer_dir())
+    ap.add_argument(
+        "--curation-target",
+        type=int,
+        default=None,
+        help="Arm B curated pair count (default: the pinned 12,000)",
+    )
+    ap.add_argument(
+        "--max-len",
+        type=int,
+        default=2048,
+        help="build cutoff, deliberately above the 1024 training "
+        "seq_len: see PinnedYields.max_len for the measured "
+        "length distribution behind this choice",
+    )
+    ap.add_argument("--held-out-frac", type=float, default=0.05)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument(
+        "--strict-pins",
+        default=True,
+        action=argparse.BooleanOptionalAction,
+        help="assert every measured count against its pinned "
+        "value. --no-strict-pins is RE-MEASUREMENT MODE: it "
+        "stamps pins_verified=false into the manifest and "
+        "sections 11/12 refuse to launch against it.",
+    )
+    ap.add_argument(
+        "--strict-host-check",
+        default=True,
+        action=argparse.BooleanOptionalAction,
+        help="disable ONLY with a deliberate reason; the check is "
+        "the guard that did not exist on 2026-08-12",
+    )
+    args = ap.parse_args()
+
+    apply_resource_guards(
+        cpu_threads=8, heavy=True, label="build_data", strict_host_check=args.strict_host_check
+    )
+
+    from align.data.build import PinnedYields, build_all
+    from align.data.sources import DataSources
+    from align.tokenizer import load_tokenizer
+
+    tokenizer = load_tokenizer(args.tokenizer)
+    sources = DataSources.from_local_hub()
+
+    kwargs = {
+        "max_len": args.max_len,
+        "held_out_frac": args.held_out_frac,
+        "seed": args.seed,
+        "strict": args.strict_pins,
+    }
+    if args.curation_target is not None:
+        kwargs["arm_b_curated"] = args.curation_target
+    pins = PinnedYields(**kwargs)
+
+    manifest = build_all(sources, args.out, tokenizer, pins=pins)
+    print(json.dumps(manifest, indent=1))
+    logger.info("build complete: %s", Path(args.out) / "manifest.json")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

code/scripts/build_stageb_data.py

@@ -0,0 +1,169 @@
+"""Build the stage-b GRPO prompt sets: keep-band filter, easy-first bands, decontamination.
+
+Consumes Big-Math (primary) and DeepMath (comparison arm) via the pinned
+streaming readers; emits per-band prompt JSONL under data/stageb/ plus a
+manifest receipt. Prompts only -- GRPO tokenizes at rollout time; there is
+no pair/tokenization build here.
+
+Decontamination is CHECKED, not assumed: Big-Math's card claims benchmark
+decontamination, but it contains GSM8K-train by construction and our
+primary eval is GSM8K-test -- so every kept problem is joined (normalized
+text) against GSM8K-test and MATH-500 and overlaps are EXCLUDED AND
+COUNTED. An eval problem inside the RL prompt pool would corrupt the
+primary contrast in the direction of the hypothesis, which is the one
+direction a defect must never point.
+
+    python scripts/build_stageb_data.py \
+        --keep-band 0.3 0.95 --band-edges 0.3 0.5 --out data/stageb
+"""
+
+import argparse
+import hashlib
+import json
+import os
+import sys
+from pathlib import Path
+
+# The package is not installed in the venv; without this a bare `python scripts/build_stageb_data.py` (the documented
+# launch line) dies with ModuleNotFoundError unless PYTHONPATH is exported (same fix as scripts/sft.py).
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from align.data.sources import DataSources, normalize_problem, read_gsm8k_test, read_math500
+from align.data.stageb_sources import (
+    BIGMATH_PIN,
+    DEEPMATH_PIN,
+    StageBSources,
+    assign_difficulty_bands,
+    iter_bigmath,
+    iter_deepmath,
+)
+
+STAGEB_HUB = "/apps/development/rlhf/hf_cache/hub"
+
+
+def _eval_text_index(sources: DataSources) -> dict[str, str]:
+    idx: dict[str, str] = {}
+    for row in read_gsm8k_test(sources):
+        idx[normalize_problem(row["question"])] = "gsm8k_test"
+    for row in read_math500(sources):
+        idx[normalize_problem(row["problem"])] = "math500"
+    return idx
+
+
+def _write_bands(name, bands, manifest, out_dir, problem_key, answer_key):
+    os.makedirs(out_dir, exist_ok=True)
+    counts = []
+    for i, band in enumerate(bands):
+        path = os.path.join(out_dir, f"{name}_band{i}.jsonl")
+        with open(path + ".partial", "w") as f:
+            for row in band:
+                f.write(
+                    json.dumps(
+                        {
+                            "problem": row[problem_key],
+                            "answer": str(row[answer_key]),
+                            "difficulty": row.get("difficulty"),
+                            "band": i,
+                            "source": row.get("source") or name,
+                        }
+                    )
+                    + "\n"
+                )
+        os.replace(path + ".partial", path)
+        with open(path, "rb") as f:
+            sha = hashlib.sha256(f.read()).hexdigest()[:16]
+        counts.append({"band": i, "n": len(band), "file": path, "sha256_16": sha})
+    manifest[name] = {"bands": counts, **manifest.get(name, {})}
+
+
+def main() -> int:
+    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    p.add_argument(
+        "--keep-band",
+        nargs=2,
+        type=float,
+        default=[0.3, 0.95],
+        metavar=("LO", "HI"),
+        help="solve-rate keep band [LO, HI) -- Big-Math only",
+    )
+    p.add_argument(
+        "--band-edges",
+        nargs="+",
+        type=float,
+        default=[0.3, 0.5],
+        help="difficulty (=1-solve_rate) edges WITHIN the keep band, easy band first",
+    )
+    p.add_argument(
+        "--deepmath-edges",
+        nargs="+",
+        type=float,
+        default=[4.0, 6.5],
+        help="DeepMath native-difficulty edges",
+    )
+    p.add_argument("--out", default="data/stageb")
+    args = p.parse_args()
+
+    sources = DataSources.from_local_hub()
+    eval_idx = _eval_text_index(sources)
+    sb = StageBSources.from_local_hub(hub=STAGEB_HUB, need_openmath=False, need_bigmath=True)
+    manifest: dict = {
+        "pins": {"bigmath": BIGMATH_PIN.revision, "deepmath": DEEPMATH_PIN.revision},
+        "keep_band_solve_rate": args.keep_band,
+        "band_edges_difficulty": args.band_edges,
+        "deepmath_edges": args.deepmath_edges,
+    }
+
+    # ---- Big-Math: keep-band + decontamination + banding
+    lo, hi = args.keep_band
+    kept, dropped_band, contaminated = [], 0, {"gsm8k_test": 0, "math500": 0}
+    for row in iter_bigmath(sb):
+        d = row.get("difficulty")
+        if d is None:
+            continue  # 4 unparseable rows, counted in the banding manifest already
+        sr = 1.0 - d
+        if not (lo <= sr < hi):
+            dropped_band += 1
+            continue
+        hit = eval_idx.get(normalize_problem(row["problem"]))
+        if hit:
+            contaminated[hit] += 1
+            continue
+        kept.append(row)
+    bands, band_manifest = assign_difficulty_bands(iter(kept), edges=args.band_edges)
+    manifest["bigmath"] = {
+        "kept": len(kept),
+        "dropped_outside_keep_band": dropped_band,
+        "excluded_eval_overlap": contaminated,
+        "banding": band_manifest.as_log_dict(),
+    }
+    _write_bands("bigmath", bands, manifest, args.out, "problem", "answer")
+
+    # ---- DeepMath comparison arm: banding on native difficulty + decontamination
+    dm_kept, dm_contam = [], {"gsm8k_test": 0, "math500": 0}
+    for row in iter_deepmath(sb):
+        hit = eval_idx.get(normalize_problem(row["question"]))
+        if hit:
+            dm_contam[hit] += 1
+            continue
+        dm_kept.append(row)
+    dm_bands, dm_manifest = assign_difficulty_bands(iter(dm_kept), edges=args.deepmath_edges)
+    manifest["deepmath"] = {
+        "kept": len(dm_kept),
+        "excluded_eval_overlap": dm_contam,
+        "banding": dm_manifest.as_log_dict(),
+    }
+    _write_bands("deepmath", dm_bands, manifest, args.out, "question", "final_answer")
+
+    mpath = os.path.join(args.out, "manifest.json")
+    with open(mpath + ".partial", "w") as f:
+        json.dump(manifest, f, indent=1)
+    os.replace(mpath + ".partial", mpath)
+    print(json.dumps({k: v for k, v in manifest.items() if k in ("bigmath", "deepmath")}, indent=1))
+    print(f"manifest: {mpath}")
+    return 0
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(main())

code/scripts/ceiling_test.py

@@ -0,0 +1,152 @@
+"""Stage-b reward-path ceiling test: a PERFECT policy must score ~100%.
+
+House rule from a paid-for incident (mistral-prm-ppo, 2026-08-14): two PPO
+runs measured a broken answer-extractor -- a perfect policy scored 0/100.
+No GRPO rollout launches until this driver certifies the EXACT reward path
+(extract_answer -> OutcomeVerifier -> CompositeReward + format instrument)
+at >= --floor on every gold source, with every miss listed in the receipt
+for eyeballing.
+
+CPU-only, no model: each (problem, gold) becomes
+`build_perfect_completion(gold, eos)` -- byte-identical to what a policy
+flawlessly imitating its SFT targets emits, EOS included (rollout decodes
+with skip_special_tokens=False).
+
+Sources: gsm8k (1,319), math500 (500), deepmath (--deepmath-n sampled
+evenly across shards; stage-b's actual gold distribution -- the reason this
+must re-run here and not only on the stage-a sets).
+
+    python scripts/ceiling_test.py --deepmath-n 5000
+"""
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+# Run from a checkout without an install: put the repo root on sys.path first.
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from align.data.schema import extract_answer, register_answer_suffixes
+from align.data.sources import DataSources, gsm8k_gold_answer, read_gsm8k_test, read_math500
+from align.data.stageb_sources import StageBSources, iter_bigmath, iter_deepmath
+from align.eval.roundtrip import build_perfect_completion
+from align.eval.verifier import OutcomeVerifier, VerifierConfig
+from align.rl.format_instrument import format_scores
+from align.rl.rewards import CompositeReward
+
+EOS = "</s>"
+
+
+#: The corpora live in two hubs: the stage-a pins under sources.DEFAULT_HUB
+#: (the voxtral cache this project inherited), the stage-b downloads under
+#: the rlhf hf_cache. Resolved per-source-set, not via one env var.
+STAGEB_HUB = "/apps/development/rlhf/hf_cache/hub"
+
+
+def _gold_sources(args) -> dict[str, list[str]]:
+    out: dict[str, list[str]] = {}
+    sources = DataSources.from_local_hub()  # DEFAULT_HUB: gsm8k/math500 pins
+    if args.gsm8k:
+        out["gsm8k"] = [gsm8k_gold_answer(r) for r in read_gsm8k_test(sources)]
+    if args.math500:
+        out["math500"] = [r["answer"] for r in read_math500(sources)]
+    if args.deepmath_n or args.bigmath_n:
+        sb = StageBSources.from_local_hub(
+            hub=args.stageb_hub,
+            need_deepmath=bool(args.deepmath_n),
+            need_openmath=False,
+            need_bigmath=bool(args.bigmath_n),
+        )
+        for key, it, total, n, field in (
+            ("deepmath", iter_deepmath, 103_110, args.deepmath_n, "final_answer"),
+            ("bigmath", iter_bigmath, 251_122, args.bigmath_n, "answer"),
+        ):
+            if not n:
+                continue
+            golds: list[str] = []
+            for i, row in enumerate(it(sb)):
+                # even sample across the stream without materializing the corpus
+                if i % max(1, total // n) == 0:
+                    golds.append(str(row[field]))
+                if len(golds) >= n:
+                    break
+            out[key] = golds
+    return out
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument("--floor", type=float, default=0.98)
+    parser.add_argument("--gsm8k", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--math500", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--deepmath-n", type=int, default=5000)
+    parser.add_argument("--bigmath-n", type=int, default=5000)
+    parser.add_argument(
+        "--workers", type=int, default=4, help="verifier procs; keep <=4 on the shared box"
+    )
+    parser.add_argument("--out", default="outputs/eval/ceiling_test.json")
+    parser.add_argument("--stageb-hub", default=STAGEB_HUB)
+    args = parser.parse_args()
+
+    register_answer_suffixes(EOS)  # BEFORE any extraction (golden.py's incident)
+    verifier = OutcomeVerifier(VerifierConfig(workers=args.workers))
+    reward = CompositeReward(format_weight=0.1)  # overlong off: length is not under test
+
+    report: dict = {"floor": args.floor, "eos": EOS, "sources": {}}
+    overall_pass = True
+    for name, golds in _gold_sources(args).items():
+        completions = [build_perfect_completion(g, EOS) for g in golds]
+        preds = [extract_answer(c) for c in completions]
+        extract_ok = sum(1 for p in preds if p is not None)
+        outcomes = verifier.verify_batch(golds, preds)
+        correct = [o.name == "CORRECT" for o in outcomes]
+        fscores = [format_scores(c)["format/score"] for c in completions]
+        rewards = [
+            reward(accuracy=1.0 if ok else 0.0, format_score=f, response_len=0)["reward/total"]
+            for ok, f in zip(correct, fscores)
+        ]
+        ceiling = sum(correct) / len(golds) if golds else float("nan")
+        misses = [
+            {"gold": g, "pred": p, "outcome": o.name}
+            for g, p, o, ok in zip(golds, preds, outcomes, correct)
+            if not ok
+        ]
+        passed = ceiling >= args.floor
+        overall_pass &= passed
+        report["sources"][name] = {
+            "n": len(golds),
+            "extract_ok_rate": extract_ok / len(golds),
+            "verify_correct_rate": ceiling,
+            "mean_format_score": sum(fscores) / len(fscores),
+            "mean_reward_total": sum(rewards) / len(rewards),
+            "passed": passed,
+            "misses": misses[:50],
+            "misses_total": len(misses),
+        }
+        print(
+            f"{name}: n={len(golds)} extract_ok={extract_ok / len(golds):.4f} "
+            f"ceiling={ceiling:.4f} misses={len(misses)} -> {'PASS' if passed else 'FAIL'}",
+            flush=True,
+        )
+
+    report["overall_pass"] = bool(overall_pass)
+    os.makedirs(os.path.dirname(args.out), exist_ok=True)
+    tmp = args.out + ".partial"
+    with open(tmp, "w") as f:
+        json.dump(report, f, indent=1)
+    os.replace(tmp, args.out)
+    print(f"receipt: {args.out} | overall_pass={overall_pass}")
+    if not overall_pass:
+        print(
+            "CEILING FAIL: eyeball every miss in the receipt before touching the "
+            "verifier -- a representational-equivalence gap and a broken extractor "
+            "look identical in the rate and completely different in the misses.",
+            file=sys.stderr,
+        )
+    return 0 if overall_pass else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

code/scripts/ckpt_convert.py

@@ -0,0 +1,142 @@
+#!/usr/bin/env python
+"""Sharded checkpoint tools: inspect, verify, export to HF safetensors, reshard.
+
+    python scripts/ckpt_convert.py inspect RUN/checkpoints/micro_N
+    python scripts/ckpt_convert.py verify  RUN/checkpoints/micro_N
+    python scripts/ckpt_convert.py export  RUN/checkpoints/micro_N OUT_DIR
+    python scripts/ckpt_convert.py reshard RUN/checkpoints/micro_N NEW_DIR --world-size 2
+
+CPU only, no GPU, no process group, no heavy-job lock: `export` is a streaming
+byte copy (peak RSS ~ a few MiB plus the Python runtime, independent of model
+size) and `reshard` re-cuts the weights the same way. The ONE part that holds
+real memory is the optimizer re-cut inside `reshard` (one NEW rank's optimizer
+state, e.g. 12 B/param/new_ws for fp32-master AdamW); it prints that estimate
+and refuses to start when the box cannot afford it (CLAUDE.md: 45 GB
+MemAvailable floor for anything heavy, no swap).
+
+WHY THIS EXISTS. `train.checkpoint_format = "sharded"` writes per-rank files
+that eval, the rollout weight sync and external tools cannot read. `export`
+turns one into exactly the directory the gathered writer would have produced
+(byte-identical; tested by file hash) without ever holding the model in RAM --
+the gathered path's rank-0 gather measured +28 GB RSS on OLMoE (docs/
+MEASUREMENTS.md, Track M). `reshard` is the ONLY way to resume at a different
+world size: resume refuses it, because sharded weights are rank-local row
+slices and silently loading them under another partitioning trains plausibly
+on scrambled weights.
+
+Exit status: 0 ok, 1 refused/corrupt (message on stderr), 2 usage.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+
+os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")  # never touch a GPU from a converter
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from align.train import sharded_ckpt as sc  # noqa: E402
+
+
+def _mem_available_gib() -> float:
+    with open("/proc/meminfo") as f:
+        for line in f:
+            if line.startswith("MemAvailable:"):
+                return int(line.split()[1]) / 2**20
+    return 0.0
+
+
+def _optimizer_bytes(src: str, ws: int) -> int:
+    return sum(os.path.getsize(os.path.join(src, f"optimizer_rank{r}.pt")) for r in range(ws) if os.path.exists(os.path.join(src, f"optimizer_rank{r}.pt")))
+
+
+def cmd_inspect(a) -> int:
+    m = sc.read_manifest(a.src)
+    total = sum(f["bytes"] for f in m["files"].values())
+    print(
+        json.dumps(
+            {
+                "world_size": m["world_size"],
+                "tensors": len(m["tensors"]),
+                "weight_bytes": total,
+                "export_family": (m.get("export") or {}).get("family"),
+                "resharded_from": m.get("resharded_from"),
+                "problem": sc.sharded_problem(a.src, lambda p: None),
+            },
+            indent=1,
+        )
+    )
+    return 0
+
+
+def cmd_verify(a) -> int:
+    m = sc.read_manifest(a.src)
+    bad = []
+    for r in range(m["world_size"]):
+        p = sc.verify_rank_file(a.src, r)
+        print(f"rank {r}: {'OK' if p is None else p}")
+        if p:
+            bad.append(p)
+    return 1 if bad else 0
+
+
+def cmd_export(a) -> int:
+    from align.utils.guard import assert_paths_off_root
+
+    assert_paths_off_root([os.path.dirname(os.path.abspath(a.dst))])
+    res = sc.export_hf(a.src, a.dst, verify=not a.no_verify)
+    print(json.dumps({**res, "peak_rss_kb": sc.peak_rss_kb()}))
+    return 0
+
+
+def cmd_reshard(a) -> int:
+    from align.utils.guard import assert_paths_off_root
+
+    assert_paths_off_root([os.path.dirname(os.path.abspath(a.dst))])
+    m = sc.read_manifest(a.src)
+    if not a.drop_optimizer:
+        est = _optimizer_bytes(a.src, m["world_size"]) / max(1, a.world_size) / 2**30
+        avail = _mem_available_gib()
+        print(f"optimizer re-cut peak ~ {est:.1f} GiB; MemAvailable {avail:.0f} GiB", file=sys.stderr)
+        if not a.allow_large and avail < est + 45:
+            print(
+                f"refusing: need MemAvailable >= {est + 45:.0f} GiB (estimate + the 45 GiB box floor) and have "
+                f"{avail:.0f}. Wait for the box, or pass --allow-large if you accept the risk, or --drop-optimizer "
+                "(fresh moments; NOT lossless).",
+                file=sys.stderr,
+            )
+            return 1
+    res = sc.reshard(a.src, a.dst, a.world_size, drop_optimizer=a.drop_optimizer)
+    print(json.dumps({**res, "peak_rss_kb": sc.peak_rss_kb()}))
+    return 0
+
+
+def main(argv=None) -> int:
+    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    sub = ap.add_subparsers(dest="cmd", required=True)
+    for name in ("inspect", "verify"):
+        p = sub.add_parser(name)
+        p.add_argument("src")
+    p = sub.add_parser("export")
+    p.add_argument("src")
+    p.add_argument("dst")
+    p.add_argument("--no-verify", action="store_true", help="skip the per-file sha256 pass (not recommended)")
+    p = sub.add_parser("reshard")
+    p.add_argument("src")
+    p.add_argument("dst")
+    p.add_argument("--world-size", type=int, required=True)
+    p.add_argument("--drop-optimizer", action="store_true", help="omit optimizer files: resume starts FRESH moments (NOT lossless)")
+    p.add_argument("--allow-large", action="store_true")
+    a = ap.parse_args(argv)
+    try:
+        return {"inspect": cmd_inspect, "verify": cmd_verify, "export": cmd_export, "reshard": cmd_reshard}[a.cmd](a)
+    except (sc.ShardedCheckpointError, FileExistsError, FileNotFoundError, KeyError) as e:
+        print(f"{type(e).__name__}: {e}", file=sys.stderr)
+        return 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

code/scripts/dedup_arm_b_shards.py

@@ -0,0 +1,308 @@
+"""One-shot, surgical dedup of the Arm B preference shards.
+
+WHY THIS EXISTS. The ref-logprob cache build refused the Arm B shards with
+`ValueError: duplicate pair_id` -- correctly. The train split held 11,404
+rows with 87 pair_ids present exactly twice, and held_out held 596 rows with
+2 doubled ids; every second occurrence is byte-identical to its first
+(same input_ids, boundaries, rejected side, and extras). Root cause: PRM800K
+carries multiple annotation records per problem, so two records can yield the
+IDENTICAL pair site; `PairCapper.offer` caps per (problem_id, step_index) but
+never checks pair_id, and `curate()` merges everything per problem
+(nmo/data/prm800k_pairs.py, `by_problem.setdefault(...).append(p)`) without a
+seen-pair_id check -- so the depth-round top-up re-selects a pair the first
+round already took.
+
+WHAT IT DOES. For each split (train, held_out):
+  1. read the live shards via the store's own reader;
+  2. keep the FIRST occurrence of every pair_id, drop second-and-later
+     occurrences -- REFUSING unless the dropped row is content-identical to
+     the kept one (a non-identical duplicate is a different bug and must not
+     be silently resolved by row order);
+  3. write fresh shards through the same ShardWriter/receipt API the build
+     used, preserving shard naming (`arm_b-train-NNNN` etc.).
+Then it asserts zero train/held pair_id overlap, atomically swaps
+`arm_b -> arm_b_predup` and `arm_b_dedup -> arm_b` (the predup directory is
+KEPT as evidence), and updates data/manifest.json in place:
+`counts.arm_b_curated`, the two `arm_b-*` shard hashes (same
+content_hash-of-receipt-hashes chain `_ShardSet.finalize` uses), and the
+`length_bias_baseline` Arm B entries re-measured over the deduped TRAIN split
+with the exact function scripts/dpo.py uses
+(align.train.dpo_stage.measure_chosen_longer_fraction). A provenance block
+`arm_b_dedup` records old -> new for everything it touched.
+
+Deterministic (pure function of the input shards; no RNG), CPU-only.
+The ref-logprob caches keyed on the old manifest hashes are stale after this
+and must be re-precomputed -- that is the point.
+
+Run:
+
+    cd /apps/development/rlhf/nano-mistral-o1
+    CUDA_VISIBLE_DEVICES= .venv/bin/python scripts/dedup_arm_b_shards.py --dry-run
+    CUDA_VISIBLE_DEVICES= .venv/bin/python scripts/dedup_arm_b_shards.py
+
+Refuses to run twice: an existing `arm_b_predup` means the swap already
+happened.
+"""
+
+from __future__ import annotations
+
+import os
+
+# CPU-only by construction: never let a lazily-imported torch see a GPU.
+os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
+
+import argparse
+import json
+import logging
+from pathlib import Path
+
+import numpy as np
+
+logger = logging.getLogger("dedup_arm_b")
+
+REPO = Path("/apps/development/rlhf/nano-mistral-o1")
+SPLITS = (("train", "arm_b-train"), ("held_out", "arm_b-held"))
+#: manifest hash key per split prefix, matching build_all's `hashes` dict.
+MANIFEST_HASH_KEY = {"arm_b-train": "arm_b-train", "arm_b-held": "arm_b-held"}
+
+
+def rows_identical(a, b) -> bool:
+    """Full content identity between two store Rows -- every stored field.
+
+    Anything short of this (say, comparing input_ids only) would let a real
+    divergence hide behind a shared pair_id and get resolved by row order.
+    """
+    return (
+        np.array_equal(a.input_ids, b.input_ids)
+        and np.array_equal(a.boundary_positions, b.boundary_positions)
+        and a.problem_id == b.problem_id
+        and a.subject == b.subject
+        and sorted(a.ragged) == sorted(b.ragged)
+        and all(np.array_equal(a.ragged[f], b.ragged[f]) for f in a.ragged)
+        and a.extras == b.extras
+    )
+
+
+def dedup_row_indices(ds) -> tuple[list[int], list[int]]:
+    """(kept, dropped) row indices: first occurrence per pair_id wins.
+
+    Raises if any dropped row is NOT content-identical to its kept first
+    occurrence -- that would be a semantic collision, not a re-added pair,
+    and this script is only licensed to remove exact re-adds.
+    """
+    first_by_pair: dict[str, int] = {}
+    kept: list[int] = []
+    dropped: list[int] = []
+    for i in range(len(ds)):
+        row = ds[i]
+        pair_id = row.extras["pair_id"]
+        if pair_id not in first_by_pair:
+            first_by_pair[pair_id] = i
+            kept.append(i)
+            continue
+        j = first_by_pair[pair_id]
+        if not rows_identical(ds[j], row):
+            raise ValueError(
+                f"pair_id {pair_id!r} occurs at rows {j} and {i} with "
+                "DIFFERENT content. That is a hash/semantics collision, not a "
+                "duplicate re-add; refusing to pick a winner by row order."
+            )
+        dropped.append(i)
+    return kept, dropped
+
+
+def write_split(ds, kept: list[int], out_dir: Path, prefix: str) -> str:
+    """Write the kept rows through the build's own shard-set API.
+
+    Returns the manifest-level content hash (content_hash over the per-shard
+    receipt hashes -- the same chain `_ShardSet.finalize` produces, so the
+    updated manifest entry is computed by the code that defined it).
+    """
+    from align.data.build import _PREF_FIELDS, _ShardSet
+
+    shard_set = _ShardSet(out_dir, prefix, _PREF_FIELDS)
+    for i in kept:
+        row = ds[i]
+        shard_set.add(
+            input_ids=row.input_ids,
+            problem_id=row.problem_id,
+            subject=row.subject,
+            boundary_positions=row.boundary_positions,
+            ragged={
+                "rejected_input_ids": row.ragged["rejected_input_ids"],
+                "rejected_boundary_positions": row.ragged["rejected_boundary_positions"],
+            },
+            **row.extras,
+        )
+    return shard_set.finalize()
+
+
+def step_basis_fraction(rows) -> float:
+    """Arm B's step-basis chosen-longer fraction over collator-shaped rows.
+
+    Mirrors build.py's `chosen_longer_fraction_step_tokens`: length measured
+    from the SHARED PREFIX (the one step that differs), counted over every
+    row. Kept next to the completion-basis measurement so the manifest's two
+    Arm B baselines stay siblings of the same dataset state.
+    """
+    longer = sum(
+        int(
+            len(r["chosen_input_ids"]) - int(r["shared_prefix_len"])
+            > len(r["rejected_input_ids"]) - int(r["shared_prefix_len"])
+        )
+        for r in rows
+    )
+    return longer / len(rows) if rows else float("nan")
+
+
+def main(argv=None) -> int:
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
+    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    ap.add_argument("--root", type=Path, default=REPO, help="repo root")
+    ap.add_argument(
+        "--dry-run", action="store_true", help="measure and report; write and swap nothing"
+    )
+    args = ap.parse_args(argv)
+
+    from align.data.loader import preference_dataset
+    from align.data.store import ShardDataset
+
+    shards_root = args.root / "data" / "shards"
+    live = shards_root / "arm_b"
+    staged = shards_root / "arm_b_dedup"
+    predup = shards_root / "arm_b_predup"
+    manifest_path = args.root / "data" / "manifest.json"
+
+    if predup.exists():
+        raise SystemExit(
+            f"{predup} already exists: the swap already happened. This script "
+            "is one-shot; a second run would dedup the already-deduped shards "
+            "and overwrite the evidence."
+        )
+    if staged.exists():
+        raise SystemExit(
+            f"{staged} already exists (a previous run died mid-write?). "
+            "Inspect and delete it before re-running; a half-written shard "
+            "set masquerades as resumable."
+        )
+
+    # -- pass 1: dedup into the staging directory ---------------------------
+    pair_ids: dict[str, set[str]] = {}
+    before_after: dict[str, tuple[int, int]] = {}
+    new_hashes: dict[str, str] = {}
+    for split, prefix in SPLITS:
+        ds = ShardDataset(live / split)
+        kept, dropped = dedup_row_indices(ds)
+        before_after[split] = (len(ds), len(kept))
+        pair_ids[split] = {ds[i].extras["pair_id"] for i in kept}
+        logger.info(
+            "%s: %d rows -> %d kept (%d exact-duplicate re-adds dropped)",
+            split,
+            len(ds),
+            len(kept),
+            len(dropped),
+        )
+        if not args.dry_run:
+            new_hashes[prefix] = write_split(ds, kept, staged / split, prefix)
+
+    overlap = pair_ids["train"] & pair_ids["held_out"]
+    if overlap:
+        raise SystemExit(
+            f"{len(overlap)} pair_ids appear in BOTH deduped splits (e.g. "
+            f"{sorted(overlap)[:3]}); the train/held split is broken and a "
+            "dedup must not paper over it."
+        )
+
+    if args.dry_run:
+        logger.info("dry run: nothing written, nothing swapped")
+        return 0
+
+    # -- pass 2: verify the staged shards through a fresh reader ------------
+    for split, _prefix in SPLITS:
+        staged_ds = ShardDataset(staged / split)
+        staged_ids = [staged_ds[i].extras["pair_id"] for i in range(len(staged_ds))]
+        if len(staged_ids) != len(set(staged_ids)) or len(staged_ids) != before_after[split][1]:
+            raise SystemExit(
+                f"staged {split} re-read: {len(staged_ids)} rows, "
+                f"{len(set(staged_ids))} unique -- the write did not produce "
+                "what the dedup computed; refusing to swap."
+            )
+
+    # -- pass 3: atomic swap; the pre-dedup shards are KEPT as evidence -----
+    live.rename(predup)
+    staged.rename(live)
+    logger.info("swapped: %s -> %s, %s -> %s", "arm_b", "arm_b_predup", "arm_b_dedup", "arm_b")
+
+    # -- pass 4: manifest update, through the numbers' own definitions ------
+    manifest = json.loads(manifest_path.read_text())
+    old = {
+        "counts.arm_b_curated": manifest["counts"]["arm_b_curated"],
+        "hashes.arm_b-train": manifest["hashes"]["arm_b-train"],
+        "hashes.arm_b-held": manifest["hashes"]["arm_b-held"],
+        "arm_b_chosen_longer_fraction": manifest["length_bias_baseline"][
+            "arm_b_chosen_longer_fraction"
+        ],
+        "arm_b_chosen_longer_fraction_step_basis": manifest["length_bias_baseline"][
+            "arm_b_chosen_longer_fraction_step_basis"
+        ],
+    }
+
+    # Re-measure the length-bias baseline over the deduped TRAIN split with
+    # the exact function scripts/dpo.py falls back to, on the exact rows it
+    # would build (read-only imports; nothing here re-implements the metric).
+    from align.train.dpo_stage import measure_chosen_longer_fraction
+
+    rows = preference_dataset(live / "train")
+    new_fraction = measure_chosen_longer_fraction(rows)
+    new_fraction_step = step_basis_fraction(rows)
+
+    n_train, n_held = before_after["train"][1], before_after["held_out"][1]
+    manifest["counts"]["arm_b_curated"] = n_train + n_held
+    manifest["hashes"]["arm_b-train"] = new_hashes["arm_b-train"]
+    manifest["hashes"]["arm_b-held"] = new_hashes["arm_b-held"]
+    manifest["length_bias_baseline"]["arm_b_chosen_longer_fraction"] = new_fraction
+    manifest["length_bias_baseline"]["arm_b_chosen_longer_fraction_step_basis"] = new_fraction_step
+    manifest["arm_b_dedup"] = {
+        "reason": (
+            "87 train + 2 held_out pair_ids were present exactly twice, every "
+            "second occurrence content-identical to its first (curation "
+            "depth-round top-up re-added pairs already selected; the merge in "
+            "prm800k_pairs.curate never deduped by pair_id). The ref-logprob "
+            "cache build correctly refused the duplicates."
+        ),
+        "script": "scripts/dedup_arm_b_shards.py",
+        "train_rows": {"before": before_after["train"][0], "after": n_train},
+        "held_out_rows": {"before": before_after["held_out"][0], "after": n_held},
+        "old": old,
+        "predup_shards": str(predup),
+    }
+    manifest_path.write_text(json.dumps(manifest, indent=1))
+
+    logger.info("train: %d -> %d rows", before_after["train"][0], n_train)
+    logger.info("held_out: %d -> %d rows", before_after["held_out"][0], n_held)
+    logger.info("counts.arm_b_curated: %s -> %d", old["counts.arm_b_curated"], n_train + n_held)
+    logger.info(
+        "arm_b_chosen_longer_fraction: %s -> %.6f (step basis %s -> %.6f)",
+        old["arm_b_chosen_longer_fraction"],
+        new_fraction,
+        old["arm_b_chosen_longer_fraction_step_basis"],
+        new_fraction_step,
+    )
+    logger.info(
+        "hashes: arm_b-train %s -> %s, arm_b-held %s -> %s",
+        old["hashes.arm_b-train"],
+        new_hashes["arm_b-train"],
+        old["hashes.arm_b-held"],
+        new_hashes["arm_b-held"],
+    )
+    logger.info(
+        "NOTE: ref caches keyed on the old manifest hashes are now stale; "
+        "re-run scripts/precompute_ref_logprobs.py before scripts/dpo.py."
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    # House rule: everything above is import-safe because spawn-based worker
+    # pools re-import __main__; the entry point stays behind this guard.
+    raise SystemExit(main())

code/scripts/dpo.py

@@ -0,0 +1,590 @@
+#!/usr/bin/env python
+"""DPO stage entry: gates -> banner -> train -> arm result manifest.
+
+    torchrun --nproc_per_node=4 scripts/dpo.py configs/dpo_argilla.toml
+    torchrun --nproc_per_node=4 scripts/dpo.py configs/dpo_steppairs.toml \
+        --lr=5e-7
+
+One script, two pre-registered arms; the config selects everything. The
+objective is section-05's chunked `fused_dpo`, the loop is section-07's
+`Trainer`, the reference logprobs are constants read from the cache
+scripts/precompute_ref_logprobs.py wrote. What this script owns is the
+decision to start and the record of what happened.
+
+Entry gates run in `dpo_stage.GATE_ORDER` and any failure blocks everything
+after it, model load included. They are ordered by cost so an invalid stage
+dies in milliseconds rather than after putting 14.5GB onto a shared no-swap
+box:
+
+  1. section-11's SFT quality-gate receipt exists AND passed;
+  2. build-manifest counts match the config, pins were verified, Arm A is not
+     marked DESCOPED;
+  3. the reference cache's six-component key matches this run exactly;
+  4. the path-bias measurement exists and its decision was honoured;
+  5. the 20-step throughput probe (wall-clock is DERIVED from it, never
+     asserted in advance);
+  6. the RSS-aware preflight for the full-FT job class.
+
+On a tripwire fire the arm re-runs ONCE with the pre-registered response from
+docs/RESULTS.md, and that configuration IS the arm's result -- written into
+the arm manifest, which section-13 reads for arm selection and the paired
+evaluation.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import os
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from align.utils.guard import apply_resource_guards  # noqa: E402
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger("dpo")
+
+
+def _split_overrides(extra: list[str]) -> list[str]:
+    """Leftover argv (what argparse did not claim) -> dotted `key=value`
+    overrides for `dpo_stage.load_stage_config`.
+
+    `--lr=5e-7` and `lr=5e-7` are the same override. Anything without an `=`
+    is refused: argparse's own positional `nargs='*'` accepted the bare form
+    and then nothing ever read it, so an override the operator believed
+    applied trained at the TOML's value (the banner printed the TOML value,
+    not the intended one). Every override is either applied or rejected.
+    """
+    overrides = [a.lstrip("-") for a in extra]
+    bad = [a for a, o in zip(extra, overrides) if "=" not in o or not o.split("=", 1)[0]]
+    if bad:
+        raise SystemExit(
+            f"unrecognized arguments {bad}: overrides must be dotted key=value pairs "
+            "(e.g. --lr=5e-7); space-separated values are not supported"
+        )
+    return overrides
+
+
+def _parse_args(argv=None):
+    """argparse for the flags, hand-split for the dotted overrides.
+
+    `--lr=5e-7` is not an argparse option, so the declared-positional form
+    rejected the documented spelling (exit 2) while accepting the bare
+    `lr=5e-7` and then never reading it. Unclaimed arguments are now the
+    overrides: applied (struct mode, via `load_stage_config`) or refused.
+    """
+    ap = argparse.ArgumentParser(allow_abbrev=False, epilog="dotted overrides: --lr=5e-7")
+    ap.add_argument("config")
+    ap.add_argument(
+        "--resume",
+        default=os.environ.get("ALIGN_RESUME") == "1",
+        action=argparse.BooleanOptionalAction,
+        help="resume from the newest valid checkpoint in run_dir instead of rotating "
+        "it away as stale (default: fresh arm; the supervisor sets ALIGN_RESUME=1 on retries)",
+    )
+    ap.add_argument("--strict-host-check", default=True, action=argparse.BooleanOptionalAction)
+    ap.add_argument("--heavy-lock", default=True, action=argparse.BooleanOptionalAction)
+    args, extra = ap.parse_known_args(argv)
+    args.overrides = _split_overrides(extra)
+    return args
+
+
+def main() -> int:
+    args = _parse_args()
+
+    from align.config import default_full_ft_config
+
+    def _rss_cfg():
+        # Project host RSS from the optimizer this stage ACTUALLY runs.
+        # default_full_ft_config defaults to offload_adamw (104.7GB
+        # projection) while the DPO arms are pinned to adamw8bit (22GB) --
+        # the mismatch made the RSS gate refuse a job 5x smaller than its
+        # model whenever the host dipped below ~150GB free.
+        c = default_full_ft_config()
+        c.optimizer.name = "adamw8bit"
+        return c
+
+    apply_resource_guards(
+        cpu_threads=8,
+        heavy=args.heavy_lock,
+        label="dpo",
+        strict_host_check=args.strict_host_check,
+        rss_cfg=_rss_cfg(),
+    )
+
+    # BEFORE anything touches CUDA -- the duplicate-GPU class, third
+    # appearance (sft.py, precompute --path-bias, now here): without
+    # set_device(local_rank) every rank's bare "cuda" is device 0, and this
+    # stage's first real 4-rank launch put four ~5.4GB policy loads on GPU 0.
+    from align.utils.dist import init_distributed
+
+    # The run dir (NCCL flight-recorder dumps land under it, never /tmp) comes
+    # from the config, which is parsed before the process group so a bad TOML
+    # fails in milliseconds.
+    from align.train.dpo_stage import load_stage_config as _early_stage_config
+
+    init_distributed(run_dir=_early_stage_config(args.config, args.overrides).run_dir)
+
+    from align.train.stage_gates import assert_distinct_devices
+
+    assert_distinct_devices()
+
+    from align.data.loader import preference_dataset
+    from align.model import checkpoint as ckpt
+    from align.model.mistral import NmoMistralForCausalLM
+    from align.train.dpo_stage import (
+        PreferenceBatchSource,
+        apply_tripwire_response,
+        arm_result_manifest,
+        assert_launch_topology,
+        build_manifest_gate,
+        load_run_ref_caches,
+        load_stage_config,
+        make_preference_sampler,
+        measure_chosen_longer_fraction,
+        read_sft_quality_gate,
+        reduce_tripwire_alarms,
+        resolve_pair_split_dir,
+        resolve_sft_checkpoint,
+        run_entry_gates,
+        write_arm_result,
+    )
+    from align.train.ref_cache import read_path_bias_report
+    from align.utils.guard import assert_paths_off_root
+
+    cfg = load_stage_config(args.config, args.overrides)
+    assert_paths_off_root([cfg.run_dir])
+    # LAUNCH REFUSAL, before any gate: world_size x micro x grad_accum must be
+    # the pre-registered global_batch the banner and horizon are quoted in.
+    # Fed by the live process group, so a 2-rank launch of a 4-rank
+    # pre-registration dies here in milliseconds, not after the model load.
+    logger.info("launch topology ok: %s", assert_launch_topology(cfg))
+    with open(cfg.manifest_path) as f:
+        manifest = json.load(f)
+
+    state: dict = {"resume": bool(args.resume)}
+
+    # -- the gate sequence ------------------------------------------------
+    # Each gate is a closure over `state` so a later gate can use what an
+    # earlier one established -- which is the point of the ORDER: reaching
+    # gate N means gates 1..N-1 held.
+    def gate_sft(_cfg):
+        state["sft_gate"] = read_sft_quality_gate(_cfg.sft_quality_gate_path)
+        # The trainer writes <run_dir>/checkpoints/micro_NNNNNNNN with no
+        # stable `final` alias; the receipt records which one passed. Resolve
+        # HERE, at the first gate, so every later consumer (ref-cache key,
+        # model load, banner) sees one path -- never a hand-composed one.
+        ckpt_path, source = resolve_sft_checkpoint(
+            _cfg.sft_checkpoint, state["sft_gate"].get("checkpoint")
+        )
+        state["sft_checkpoint"] = ckpt_path
+        state["sft_gate"]["checkpoint"] = ckpt_path
+        return {"checkpoint": ckpt_path, "checkpoint_source": source}
+
+    def gate_manifest(_cfg):
+        result = build_manifest_gate(_cfg, manifest)
+        state["descoped"] = result["descoped"]
+        return result
+
+    def gate_ref_cache(_cfg):
+        # Content-hash key (checkpoint WEIGHT BYTES, not the tensor inventory);
+        # forward_impl is read off the cache itself, not the path-bias report --
+        # see dpo_stage.expected_ref_cache_key for why each.
+        state["ref_cache"], state["heldout_cache"] = load_run_ref_caches(
+            _cfg, manifest, state["sft_checkpoint"]
+        )
+        state["ref_cache_key"] = state["ref_cache"].key
+        return {"pairs": len(state["ref_cache"]), "heldout_pairs": len(state["heldout_cache"])}
+
+    def gate_path_bias(_cfg):
+        state["path_bias"] = read_path_bias_report(_cfg.path_bias_report_path)
+        return {"action": state["path_bias"]["action"], "ratio": state["path_bias"]["ratio"]}
+
+    def gate_probe(_cfg):
+        # Filled after the model exists; the probe is the one gate that needs
+        # the loop, so it records its intent here and the banner reads the
+        # measured value below.
+        state["probe"] = {
+            "step_time_s_median": float("nan"),
+            "projected_wall_clock_h": float("nan"),
+            "status": "PENDING",
+        }
+        return {"status": "PENDING until the 20-step probe runs"}
+
+    def gate_rss(_cfg):
+        from align.train.probes import projected_job_rss
+
+        state["rss"] = projected_job_rss(_rss_cfg())
+        return dict(state["rss"])
+
+    gate_report = run_entry_gates(
+        cfg,
+        {
+            "sft_quality_gate": gate_sft,
+            "build_manifest": gate_manifest,
+            "ref_cache_key": gate_ref_cache,
+            "path_bias": gate_path_bias,
+            "throughput_probe": gate_probe,
+            "rss_preflight": gate_rss,
+        },
+    )
+
+    if state["descoped"]:
+        # A descoped arm is a RESULT, not an absence: section-13's selection
+        # rule reads this manifest and must be able to tell "measured badly"
+        # from "never measured".
+        out = arm_result_manifest(
+            cfg=cfg, descoped=True, summary={}, heldout={}, banner={}, checkpoint=None
+        )
+        path = write_arm_result(os.path.join(cfg.run_dir, "arm_result.json"), out)
+        logger.warning("%s is DESCOPED; wrote %s and did not train", cfg.arm, path)
+        return 0
+
+    # -- data + model ------------------------------------------------------
+    # DISCOVERY, not path composition: the configs point at the real split
+    # dirs, and resolve_pair_split_dir re-finds them from the receipts the
+    # store wrote -- so a layout change moves this by itself, and a wrong
+    # configured path fails naming the splits that ARE present.
+    rows = preference_dataset(resolve_pair_split_dir(cfg.shard_dir, f"{cfg.arm}-train"))
+    heldout_rows = preference_dataset(
+        resolve_pair_split_dir(cfg.heldout_shard_dir, f"{cfg.arm}-held")
+    )
+    pad_id = int(manifest.get("pinned", {}).get("pad_token_id", 0))
+
+    baselines = manifest.get("length_bias_baseline", {})
+    baseline_key = f"{cfg.arm}_chosen_longer_fraction"
+    if baseline_key in baselines:
+        length_baseline = (baselines[baseline_key], "build_manifest")
+    else:
+        # Section-03 records only Arm A's baseline. Measure this arm's from
+        # its own shards rather than borrow the other's -- a tripwire armed
+        # against the wrong distribution fires forever or never.
+        length_baseline = (measure_chosen_longer_fraction(rows), "measured_at_launch")
+    logger.info("length-drift baseline: %.4f (%s)", *length_baseline)
+
+    model = _load_policy(
+        NmoMistralForCausalLM, ckpt, state["sft_checkpoint"], cfg.attn_implementation
+    )
+    # The RESOLVED backend, read back off the model: the constructor routed
+    # the request through the gated selector, so this is what actually runs.
+    state["attn_backend"] = model.attn_implementation
+    run_cfg = _run_config(cfg, length_baseline[0])
+
+    from align.train.fsdp import apply_activation_checkpointing, apply_fsdp
+
+    dist_model = apply_fsdp(apply_activation_checkpointing(model))
+    import torch as _torch
+
+    # Post-shard residency: a rank should hold ~3.6GB (its shard). Anything
+    # near 18GB means the pre-shard replica leaked past fully_shard.
+    logger.info("post-shard cuda allocated: %.2f GiB", _torch.cuda.memory_allocated() / 2**30)
+
+    source = PreferenceBatchSource(
+        rows,
+        make_preference_sampler(
+            rows,
+            max_tokens=run_cfg.train.seq_len * cfg.micro_batch_size,
+            max_rows=cfg.micro_batch_size,
+            seed=run_cfg.train.seed,
+        ),
+        pad_token_id=pad_id,
+    )
+
+    result = _train_arm(
+        cfg,
+        run_cfg,
+        dist_model,
+        source,
+        state,
+        heldout_rows,
+        pad_id,
+        gate_report,
+        manifest,
+        length_baseline,
+    )
+
+    # -- tripwire response -------------------------------------------------
+    # The rerun is collective, so the verdict is drawn from the UNION of every
+    # rank's alarms: rank-local alarms could send some ranks into a second
+    # model load + FSDP wrap while the rest exit.
+    alarms = reduce_tripwire_alarms(result["summary"].get("tripwire_alarms", []))
+    rerun_cfg, tripwire = apply_tripwire_response(cfg, alarms)
+    if tripwire["rerun_required"]:
+        logger.error(
+            "chosen-logprob collapse detected; re-running %s with the pre-registered response: %s",
+            cfg.arm,
+            tripwire["configuration"],
+        )
+        model = _load_policy(
+            NmoMistralForCausalLM, ckpt, state["sft_checkpoint"], cfg.attn_implementation
+        )
+        state["attn_backend"] = model.attn_implementation
+        dist_model = apply_fsdp(apply_activation_checkpointing(model))
+        result = _train_arm(
+            rerun_cfg,
+            _run_config(rerun_cfg, length_baseline[0]),
+            dist_model,
+            source,
+            state,
+            heldout_rows,
+            pad_id,
+            gate_report,
+            manifest,
+            length_baseline,
+            suffix="_response",
+        )
+        cfg = rerun_cfg
+
+    out = arm_result_manifest(
+        cfg=cfg,
+        descoped=False,
+        summary=result["summary"],
+        heldout=result["heldout"],
+        banner=result["banner"],
+        checkpoint=result["checkpoint"],
+        tripwire=tripwire,
+    )
+    out["gate_report"] = gate_report
+    path = write_arm_result(os.path.join(cfg.run_dir, "arm_result.json"), out)
+    logger.info("%s complete: %s", cfg.arm, path)
+    return 0
+
+
+def _load_policy(model_cls, ckpt, path: str, attn_implementation: str):
+    """Load the policy with the CONFIG-DRIVEN attention backend.
+
+    The string is never handed raw to a constructor in the sense that
+    matters: `NmoMistralForCausalLM.__init__` routes it through
+    `align.model.attention.select_attn_implementation`, which raises if the
+    request cannot be honoured -- for the pre-registered `triton_fa2` that
+    means section-06's equivalence-gate receipt must be valid for this kernel
+    source on this GPU. Silent degradation to sdpa is exactly the failure the
+    selector exists to prevent.
+    """
+    import torch
+
+    from align.model.hf_config import read_config
+
+    model_cfg = read_config(path)
+    model = model_cls.from_state_dict(
+        model_cfg,
+        ckpt.hf_to_nmo(ckpt.load_hf_safetensors(path)),
+        attn_implementation=attn_implementation,
+        device="cuda",
+    )
+    return model.to("cuda", dtype=torch.bfloat16)
+
+
+def _run_config(cfg, length_baseline: float):
+    """Stage config -> section-07 RunConfig. The two are different objects on
+    purpose: RunConfig is the trainer's surface and DpoStageConfig is the
+    arm's pre-registration, and collapsing them would let a trainer default
+    silently become an experimental parameter."""
+    from align.config import RunConfig
+
+    run = RunConfig()
+    # adamw8bit per the 2026-08-17 pre-registration amendment ("SFT (and the
+    # DPO arms) run adamw8bit"). Left unset, RunConfig's offload_adamw
+    # default leaked through and the pinned-staging guard refused with a
+    # 27GiB request on a no-swap box -- correctly.
+    run.optimizer.name = "adamw8bit"
+    run.optimizer.lr = cfg.lr
+    run.optimizer.grad_accum = cfg.grad_accum
+    run.train.epochs = cfg.epochs
+    run.train.micro_batch_size = cfg.micro_batch_size
+    run.train.save_every_micro_steps = cfg.save_every_micro_steps
+    run.train.keep_checkpoints = cfg.keep_checkpoints
+    # Forwarded explicitly: RunConfig is built from defaults here, so a stage
+    # field that is not copied is a silently disarmed config.
+    run.train.per_rank_optimizer_state = cfg.per_rank_optimizer_state
+    run.optimizer.max_grad_norm = cfg.max_grad_norm
+    run.train.max_wall_seconds = cfg.max_wall_seconds
+    run.train.max_consecutive_skipped_steps = cfg.max_consecutive_skipped_steps
+    run.train.max_skipped_fraction = cfg.max_skipped_fraction
+    run.train.full_ft = True
+    run.data.num_workers = cfg.num_workers
+    run.tripwires.chosen_longer_baseline = length_baseline
+    run.tripwires.length_band = cfg.length_band
+    run.tripwires.logps_window = cfg.logps_window
+    run.tripwires.armed = tuple(cfg.armed_tripwires)
+    # The tripwire may speak exactly when its estimator is full. These are two
+    # different quantities -- `length_window_pairs` is the span the adapter
+    # averages over, `length_min_pairs` is when the tripwire is allowed to
+    # fire -- and they are deliberately the SAME number: a min_pairs below the
+    # window size would let the guard speak from a partially-filled window,
+    # which is precisely the noisy regime both settings exist to avoid.
+    run.tripwires.length_min_pairs = cfg.length_window_pairs
+    run.run_name = f"dpo_{cfg.arm}"
+    return run
+
+
+def _train_arm(
+    cfg,
+    run_cfg,
+    dist_model,
+    source,
+    state,
+    heldout_rows,
+    pad_id,
+    gate_report,
+    manifest,
+    length_baseline,
+    suffix: str = "",
+):
+    """Probe -> banner -> train -> held-out accuracy. In that order, because
+    the banner quotes the PROBE-DERIVED wall-clock and the plan forbids
+    asserting a wall-clock in advance."""
+    from align.data.loader import collate_preference
+    from align.model.attention import banner_value
+    from align.train.dpo_stage import (
+        DpoBatch,
+        dpo_banner,
+        make_dpo_stage_adapter,
+        render_dpo_banner,
+    )
+    from align.train.ref_cache import heldout_preference_accuracy
+    from align.train.trainer import Trainer
+
+    run_dir = cfg.run_dir + suffix
+    # FRESH-ARM HYGIENE: a non-resuming arm run must not share a checkpoint
+    # dir with a previous run's leftovers -- stale higher-numbered dirs
+    # shadowed the rerun's fresh saves under keep-N pruning (2026-08-20).
+    from align.train.dpo_stage import probe_micro_cap, rotate_stale_checkpoints
+
+    ckpt_dir = os.path.join(run_dir, "checkpoints")
+    # Rotation is for FRESH arms only. A resuming arm (--resume, or a
+    # supervisor retry via ALIGN_RESUME=1) must keep its checkpoints: rotating
+    # them would silently restart the arm from step 0 while the supervisor
+    # believes it is continuing -- the replay defect, at the stage level.
+    resuming = bool(state.get("resume"))
+    # Collective: rank 0 renames, every rank waits (a bare per-rank rename
+    # let one rank win and killed the other three with FileNotFoundError).
+    if not resuming:
+        rotate_stale_checkpoints(ckpt_dir)
+    adapter = make_dpo_stage_adapter(state["ref_cache"], cfg, require_fused=True)
+    trainer = Trainer(run_cfg, dist_model, source, adapter, run_dir=run_dir, stage="dpo")
+
+    # REFUSE a horizon that disagrees with the corpus, before a single step.
+    # Arm A ran a 146-step horizon over a corpus needing ~289 per-rank steps:
+    # the cosine finished at half an epoch, clamped the learning rate to
+    # exactly 0, and the loop kept stepping on frozen weights for ~24
+    # GPU-hours. The Trainer now stops AT the horizon, which turns that into a
+    # short run rather than a zombie one -- this check is the other half, and
+    # catches it before the GPU-hours rather than after.
+    from align.train.dpo_stage import assert_dpo_horizon_consistent
+    from align.utils.dist import world_size as _world_size
+
+    horizon_report = assert_dpo_horizon_consistent(
+        trainer.scheduler.total_steps,
+        n_pairs=len(source.rows),
+        world_size=_world_size(),
+        micro_batch_size=cfg.micro_batch_size,
+        grad_accum=cfg.grad_accum,
+        epochs=run_cfg.train.epochs,
+        global_batch=cfg.global_batch,
+    )
+    logger.info("horizon ok: %s", horizon_report)
+
+    if resuming:
+        resumed_from = trainer.resume_from_latest()
+        logger.info(
+            "RESUMED from %s at optimizer step %d / micro step %d",
+            resumed_from,
+            trainer.optimizer_step,
+            trainer.micro_step,
+        )
+        state["resumed_from"] = resumed_from
+
+    tokens_per_step = (
+        run_cfg.train.seq_len * cfg.micro_batch_size * 2 * cfg.grad_accum
+    )  # 2: both sides of every pair
+    # The probe is the FIRST probe_steps real micro-steps of the run itself:
+    # a capped train(), report captured, cap removed, training continues from
+    # the stream's exact position. The previous construction asked a trainer
+    # that had never stepped for its step times, and throughput_report
+    # correctly refused an empty probe. Probe steps here are real training
+    # (counted in the horizon), so nothing is spent twice; the banner prints
+    # after 20 micro-steps instead of before any. The cap is on the ABSOLUTE
+    # micro-step counter, so a resumed arm probes its first probe_steps AFTER
+    # the resume point (a bare probe_steps cap trained one micro-step and
+    # crashed throughput_probe on every supervisor retry).
+    run_cfg.train.max_micro_steps = probe_micro_cap(trainer.micro_step, run_cfg.train.probe_steps)
+    trainer.train()
+    run_cfg.train.max_micro_steps = None
+    if trainer.stop_exit_code:
+        # A coordinated stop (SIGTERM / train.max_wall_seconds) saved at an
+        # optimizer-step boundary; the stage is NOT finished. Exit with the
+        # documented status instead of evaluating / reporting a half-trained arm.
+        raise SystemExit(trainer.stop_exit_code)
+    if resuming and len(trainer.step_times) < trainer.accum:
+        # Resumed within one optimizer step of the horizon (or after the last
+        # step): the probe window cannot fill a group. The run still finishes;
+        # the banner says the wall-clock was not measured rather than raising.
+        probe = {
+            "step_time_s_median": float("nan"),
+            "projected_wall_clock_h": float("nan"),
+            "status": "UNMEASURED_ON_RESUME",
+        }
+    else:
+        probe = trainer.throughput_probe(tokens_per_step)
+    state["probe"] = probe
+
+    banner = dpo_banner(
+        cfg=cfg,
+        manifest=manifest,
+        ref_cache_key=state["ref_cache_key"],
+        path_bias=state["path_bias"],
+        probe=probe,
+        rss=state["rss"],
+        backend=banner_value(
+            # The RESOLVED backend, never a literal: FSDP2 mutates the module
+            # in place so the attribute normally survives wrapping, and the
+            # fallback is what the selector resolved at load -- a hardcoded
+            # "sdpa" here would let the banner contradict the kernel that ran.
+            dist_model.attn_implementation
+            if hasattr(dist_model, "attn_implementation")
+            else state["attn_backend"]
+        ),
+        sft_gate=state["sft_gate"],
+        length_baseline=length_baseline,
+    )
+    print(render_dpo_banner(banner))
+
+    summary = trainer.train()
+    if trainer.stop_exit_code:
+        # A coordinated stop (SIGTERM / train.max_wall_seconds) saved at an
+        # optimizer-step boundary; the stage is NOT finished. Exit with the
+        # documented status instead of evaluating / reporting a half-trained arm.
+        raise SystemExit(trainer.stop_exit_code)
+    checkpoint = trainer.save_checkpoint()
+
+    import torch as _torch
+
+    _dev = _torch.device("cuda", _torch.cuda.current_device())
+    # .to(_dev): the held-out path missed the device move the training
+    # adapter got -- CPU labels reached the CUDA gather AFTER the full arm
+    # had trained and saved (Arm A rerun, 2026-08-20).
+    heldout_batches = [
+        collate_preference(heldout_rows[i : i + cfg.micro_batch_size], pad_token_id=pad_id).to(_dev)
+        for i in range(0, len(heldout_rows), cfg.micro_batch_size)
+    ]
+    heldout = heldout_preference_accuracy(
+        dist_model, heldout_batches, state["heldout_cache"], beta=cfg.beta
+    )
+    logger.info(
+        "held-out preference accuracy: %.4f [%.4f, %.4f] n=%d",
+        heldout["accuracy"],
+        heldout["ci_low"],
+        heldout["ci_high"],
+        heldout["n"],
+    )
+    return {"summary": summary, "heldout": heldout, "banner": banner, "checkpoint": checkpoint}
+
+
+if __name__ == "__main__":
+    from align.exit_codes import run_with_exit_codes
+
+    raise SystemExit(run_with_exit_codes(main))

code/scripts/eval.py

@@ -0,0 +1,102 @@
+"""Thin eval entry: single / paired / golden.
+
+    python scripts/eval.py configs/eval.toml --mode=single --model.path=...
+
+ORDER OF OPERATIONS IS LOAD-BEARING
+-----------------------------------
+1. `apply_resource_guards` runs ONLY when we are not inside a multiprocessing
+   child. The verifier's worker pool uses 'spawn', which RE-IMPORTS this
+   module in every worker; a resource guard or an exclusive flock at module
+   scope therefore executes in each child and kills it on startup. Every
+   future then dies with BrokenProcessPool -- which is precisely the
+   documented cause of the "all answers incorrect" masquerade the verifier
+   raises on. The guard is the fix; `_in_multiprocessing_child()` is what
+   makes it safe.
+2. Banner next: resolved config, dataset hashes, MDE, model hash, backend.
+   A run that cannot answer its own question does not begin.
+3. Then dispatch.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if REPO_ROOT not in sys.path:
+    sys.path.insert(0, REPO_ROOT)
+
+from align.utils.guard import _in_multiprocessing_child  # noqa: E402
+
+# See docstring point 1. This condition is not an optimization.
+if not _in_multiprocessing_child():
+    from align.utils.guard import apply_resource_guards
+
+    apply_resource_guards(cpu_threads=8, heavy=False, label="eval")
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(description="nmo evaluation entry")
+    p.add_argument("config", nargs="?", default="configs/eval.toml")
+    p.add_argument("--mode", choices=("single", "paired", "golden"), default="single")
+    p.add_argument(
+        "--out", default=None, help="report path; defaults to outputs/eval/ under the repo"
+    )
+    p.add_argument(
+        "overrides", nargs="*", default=[], help="dotted overrides, e.g. eval.limit=8"
+    )
+    return p
+
+
+def default_out_dir() -> str:
+    """Reports land under /apps -- asserted, never assumed."""
+    from align.utils.guard import assert_paths_off_root
+
+    out = os.path.join(REPO_ROOT, "outputs", "eval")
+    os.makedirs(out, exist_ok=True)
+    assert_paths_off_root([out])
+    return out
+
+
+def main(argv=None) -> int:
+    args = build_parser().parse_args(argv)
+
+    from align.config import load_toml
+
+    # Schemaless on purpose (eval.toml has no dataclass): plain nested dict,
+    # dotted overrides TOML-literal parsed. Same shape OmegaConf.to_object gave.
+    cfg = load_toml(args.config, list(args.overrides))
+    out_dir = args.out or default_out_dir()
+
+    if args.mode == "golden":
+        from align.eval.golden import run_golden_regression
+
+        banner: dict = {}
+        path = os.path.join(out_dir, "checkpoint-200-golden.json")
+        # `run_golden_regression` writes the report ITSELF, before gating, so
+        # a failing gate still leaves the evidence. Writing again here would
+        # be dead code on the path that matters: the gate raises, and anything
+        # after this call never runs.
+        try:
+            run_golden_regression(banner=banner, out_path=path)
+        finally:
+            # U01-9: name the eval source that produced the report (also when the gate raised), so
+            # `record_golden_from_report` cannot re-certify it after align/eval/ changes.
+            if os.path.isfile(path):
+                from align.train.stage_gates import stamp_golden_report
+
+                stamp_golden_report(path)
+        print(f"golden gate PASSED; report at {path}")
+        return 0
+
+    raise SystemExit(
+        f"mode {args.mode!r} needs a generation backend, which arrives "
+        "with section-04 (nmo model) / section-09 (engine). The harness "
+        "itself is complete and tested: run_eval takes a generate_fn. "
+        "Use --mode=golden, which uses plain HF generation by design."
+    )
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

code/scripts/eval_checkpoint.py

@@ -0,0 +1,175 @@
+"""Standalone checkpoint eval: fresh process, single GPU, golden-proven generator.
+
+Promotion of the .tmp/eval_gate.py driver (task #26) into a parameterized
+entry point. The design facts it carries forward are all incident-derived:
+
+  * FRESH PROCESS, never in-trainer: the train->eval handoff inside one
+    process produced the rank race (peers built a second 7B while rank 0
+    saved) and the orphan-GPU-process incident. A checkpoint on disk plus a
+    fresh process has neither failure mode.
+  * EVERYTHING under the __main__ guard: the verifier's scoring pool uses
+    the 'spawn' context, which RE-IMPORTS __main__ in every worker -- the
+    unguarded first driver had all 8 workers load their own 7B onto the
+    eval GPU.
+  * greedy generation via `golden_generate_fn` (in-repo Engine, same encode
+    as the golden regression: BOS default, truncation 768). The anchors and
+    banked SFT rows were measured with HF generate; the generator swap is
+    certified by the golden regression + checkpoint roundtrip gate, and the
+    generator is recorded in the report config.
+  * Atomic report write (.partial -> os.replace): a half-written report
+    masquerades as a complete one.
+
+Phase-0 usage (the Arm B result side of the primary McNemar):
+
+    CUDA_VISIBLE_DEVICES=0 python scripts/eval_checkpoint.py \
+        --checkpoint runs/dpo_arm_b_response/checkpoints/micro_XXXXXXXX \
+        --name arm-b-final --gsm8k-n 1319 --math500-n 0
+
+`--quality-gate` additionally runs the SFT-stage quality gate + ledger
+records; leave it off for non-SFT checkpoints (the gate's anchors are
+SFT-specific).
+"""
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+# Run from a checkout without an install: put the repo root on sys.path first.
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from align.utils.guard import apply_resource_guards
+
+
+def quality_gate_size_errors(cfg, math500_n: int, gsm8k_n: int, per_set: dict | None = None) -> list[str]:
+    """Why a run of these sizes may NOT write the SFT quality-gate receipt (empty = fine).
+
+    The receipt unlocks DPO. `assert_meets_golden` accepts any point estimate inside the target's own
+    Wilson CI, which is enormous at tiny n (1/5 passes the 0.09 anchor), so the gate certifies only a
+    FULL-size evaluation: both sets, at the configured n the anchors were recorded at. `per_set` (the
+    measured report) is checked too, in case a set held fewer problems than requested."""
+    want = {"math500": cfg.eval.math500_n, "gsm8k": cfg.eval.gsm8k_n}
+    asked = {"math500": math500_n, "gsm8k": gsm8k_n}
+    errs = [f"{k}: requested n={asked[k]}, the quality gate needs the full configured n={want[k]}"
+            for k in want if asked[k] != want[k]]
+    for k, w in want.items():
+        got = (per_set or {}).get(k, {}).get("n") if per_set is not None else w
+        if got != w:
+            errs.append(f"{k}: measured n={got}, the quality gate needs n={w}")
+    return errs
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument("--checkpoint", required=True, help="checkpoint dir (HF-layout safetensors dir)")
+    parser.add_argument("--config", default="configs/sft.toml")
+    parser.add_argument("--name", required=True, help="report model_name and output stem")
+    parser.add_argument("--out", default=None, help="default: outputs/eval/<name>.json")
+    parser.add_argument("--math500-n", type=int, default=None, help="0 skips; default: config")
+    parser.add_argument("--gsm8k-n", type=int, default=None, help="0 skips; default: config")
+    parser.add_argument(
+        "--quality-gate", action="store_true", help="SFT-stage gate + ledger records"
+    )
+    args = parser.parse_args()
+
+    # Guards BEFORE torch import (env vars are load-bearing); heavy=True takes
+    # the exclusive job lock -- a 7B generation eval over 1,319 problems is a
+    # multi-hour GPU job and must serialise behind training like any other.
+    apply_resource_guards(cpu_threads=8, heavy=True, label=f"eval-{args.name}")
+
+
+    from align.config import load_config  # noqa: E402
+    from align.eval.golden import golden_generate_fn, load_checkpoint_model  # noqa: E402
+    from align.eval.harness import EvalConfig, load_eval_set, run_eval  # noqa: E402
+    from align.eval.verifier import OutcomeVerifier  # noqa: E402
+    from align.tokenizer import load_tokenizer  # noqa: E402
+    from align.train.stage_smoke import _manifest_hashes, _mde  # noqa: E402
+
+    cfg = load_config(args.config)
+    checkpoint = args.checkpoint
+    assert os.path.isdir(checkpoint), checkpoint
+    out_path = args.out or os.path.join("outputs", "eval", f"{args.name}.json")
+
+    math500_n = cfg.eval.math500_n if args.math500_n is None else args.math500_n
+    gsm8k_n = cfg.eval.gsm8k_n if args.gsm8k_n is None else args.gsm8k_n
+    if args.quality_gate:
+        errs = quality_gate_size_errors(cfg, math500_n, gsm8k_n)
+        if errs:
+            print("refusing --quality-gate on a subset eval (the receipt unlocks DPO):\n  " + "\n  ".join(errs),
+                  file=sys.stderr)
+            return 2
+    problems = []
+    # 0 in this CLI means SKIP the set (checked here); upstream load_eval_set
+    # uses limit=0 for "no limit", which this branch structure never passes
+    # ambiguously -- a negative value falls through to upstream validation.
+    if math500_n != 0:
+        problems += load_eval_set(cfg.paths.math500_path, limit=math500_n)
+    if gsm8k_n != 0:
+        problems += load_eval_set(cfg.paths.gsm8k_path, limit=gsm8k_n)
+    if not problems:
+        print("nothing to evaluate (both sets skipped)", file=sys.stderr)
+        return 2
+
+    # After every refusal above: a 7B load is the expensive step.
+    tokenizer = load_tokenizer(checkpoint)
+    model = load_checkpoint_model(checkpoint, device="cuda")
+
+    verifier = OutcomeVerifier()
+    gen_fn = golden_generate_fn(model, tokenizer, "cuda")
+    print(f"evaluating {args.name}: {len(problems)} problems from {checkpoint}", flush=True)
+    rep = run_eval(problems, gen_fn, verifier, EvalConfig(), model_name=args.name)
+    per_set = rep.per_set
+    for name, scores in sorted(per_set.items()):
+        print(f"  {name}: n={scores.get('n')} pass@1={scores.get('pass@1')}", flush=True)
+
+    report = {
+        "per_set": per_set,
+        "config": {
+            "mde": _mde(cfg),
+            "n_problems": sum(r.get("n", 0) for r in per_set.values()),
+            "eval_set_hash": _manifest_hashes(cfg),
+            "generator": "align Engine greedy (golden_generate_fn)",
+            "checkpoint": checkpoint,
+        },
+        "pooled": rep.pooled,
+        "per_problem": rep.per_problem,
+    }
+    os.makedirs(os.path.dirname(out_path), exist_ok=True)
+    tmp = out_path + ".partial"
+    with open(tmp, "w") as f:
+        json.dump(report, f, indent=1, default=str)
+    os.replace(tmp, out_path)
+    print(f"report written: {out_path}", flush=True)
+
+    if not args.quality_gate:
+        return 0
+
+    from align.train.stage_gates import (  # noqa: E402
+        GateLedger,
+        eval_source_hash,
+        quality_gate,
+        write_quality_gate_receipt,
+    )
+
+    errs = quality_gate_size_errors(cfg, math500_n, gsm8k_n, per_set=per_set)
+    if errs:  # report is already written above; only the receipt is withheld
+        print("refusing to write the quality-gate receipt:\n  " + "\n  ".join(errs), file=sys.stderr)
+        return 2
+
+    run_dir = os.path.join(cfg.paths.checkpoint_dir, "sft")
+    ledger = GateLedger(run_dir, eval_source_hash=eval_source_hash())
+    ledger.record("eval", passed=True, detail=json.dumps(per_set, default=str))
+    decision = quality_gate(report, cfg)
+    ledger.record("quality_gate", passed=decision.passed, detail=decision.halt_message[:400])
+    write_quality_gate_receipt(cfg, decision, checkpoint=checkpoint)
+    print(f"QUALITY GATE: passed={decision.passed}", flush=True)
+    if not decision.passed:
+        print(decision.halt_message, file=sys.stderr, flush=True)
+    return 0 if decision.passed else 1
+
+
+if __name__ == "__main__":
+    # See module docstring: spawn workers re-import __main__; nothing above
+    # this guard may build a model or touch the GPU.
+    sys.exit(main())

code/scripts/gate_flash_attention.py

@@ -0,0 +1,132 @@
+"""Run the triton_fa2 equivalence suite and, ONLY on success, write the gate.
+
+Run (pin ONE free device):
+    CUDA_VISIBLE_DEVICES=<free gpu> uv run python scripts/gate_flash_attention.py
+
+This is the mechanism behind the never-silently-claimed rule. The gate artifact
+is the ONLY thing that makes `triton_fa2` selectable
+(`align.kernels.gate.resolve_triton_fa2_backend` raises without it), and this
+script is the only writer. It records the kernel source hash, the torch/triton
+versions, the exact GPU, and the tolerances and shapes the suite covered, so a
+later edit to the kernel or a move to a different card invalidates it
+automatically.
+
+Deliberately a separate entry point rather than a pytest hook: writing the gate
+must require an explicit human action on a machine with a real GPU, and must
+never be a side effect of someone running the test suite with `-k` narrowing
+the selection to a subset.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import xml.etree.ElementTree as ET
+
+REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, REPO)
+
+from align.kernels import gate as gate_mod
+
+#: Kept in step with tests/test_flash_attention.py. Recorded in the artifact so
+#: a reader knows what "verified" actually covered.
+TOLERANCES = {
+    "policy": "max |err| vs fp32 reference <= max(3.0 x torch-SDPA-bf16 error, 2e-2)",
+    "slack_over_sdpa": 3.0,
+    "absolute_floor": 2e-2,
+}
+
+SHAPES_COVERED = [
+    "gqa (n_q_heads, n_kv_heads) in [(8,2), (32,8)] plus MHA-degenerate (4,4)",
+    "head_dim in [64, 128]",
+    "seq in [192, 256, 512, 1000] (block-divisible AND non-divisible)",
+    "causal and non-causal",
+    "key_padding_mask, right-padded, non-block-aligned",
+    "every autotune config individually pinned (fwd + both bwd kernels)",
+    "dtype bfloat16 only (fp16/fp32 rejected)",
+]
+
+
+def _count_outcomes(junit_path: str) -> tuple[int, int, int]:
+    """(passed, skipped, failed+errored) from a pytest junit report.
+
+    pytest exits 0 when EVERY test is skipped (no CUDA, more than one device
+    visible, heavy-job lock held), so the exit status alone cannot prove the
+    kernel ran. Raises if the report is missing or unparseable.
+    """
+    root = ET.parse(junit_path).getroot()
+    passed = skipped = bad = 0
+    for case in root.iter("testcase"):
+        if case.find("skipped") is not None:
+            skipped += 1
+        elif case.find("failure") is not None or case.find("error") is not None:
+            bad += 1
+        else:
+            passed += 1
+    return passed, skipped, bad
+
+
+def main() -> int:
+    cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "")
+    devices = [d.strip() for d in cvd.split(",")]
+    if len(devices) != 1 or not devices[0]:
+        print(
+            "refusing to run: set CUDA_VISIBLE_DEVICES to exactly ONE free device "
+            f"(got {cvd!r}; shared box, and the suite skips unless exactly one is visible).",
+            file=sys.stderr,
+        )
+        return 2
+
+    junit = os.path.join(REPO, ".cache", "gate_flash_attention_junit.xml")
+    os.makedirs(os.path.dirname(junit), exist_ok=True)
+    if os.path.exists(junit):
+        os.remove(junit)  # a stale report must never vouch for this run
+    cmd = [
+        sys.executable,
+        "-m",
+        "pytest",
+        os.path.join(REPO, "tests", "test_flash_attention.py"),
+        "-m",
+        "gpu",
+        "-q",
+        "-rs",
+        f"--junitxml={junit}",
+    ]
+    print("running the equivalence suite:", " ".join(cmd), flush=True)
+    result = subprocess.run(cmd, cwd=REPO)
+
+    if result.returncode != 0:
+        print(
+            "\nequivalence suite FAILED -- gate NOT written. triton_fa2 stays "
+            "unselectable; SDPA remains the backend.",
+            file=sys.stderr,
+        )
+        return result.returncode
+
+    # Exit 0 is necessary, not sufficient: an all-skipped run also exits 0.
+    try:
+        passed, skipped, bad = _count_outcomes(junit)
+    except (OSError, ET.ParseError) as e:
+        print(f"\ncannot read the junit report ({e!r}) -- gate NOT written.", file=sys.stderr)
+        return 1
+    if passed == 0 or skipped != 0 or bad != 0:
+        print(
+            f"\nequivalence suite did not EXECUTE in full (passed={passed}, "
+            f"skipped={skipped}, failed={bad}) -- gate NOT written. Skips mean "
+            "no CUDA device, a device count != 1, or the heavy-job lock held "
+            "by a live job; none of those verifies the kernel.",
+            file=sys.stderr,
+        )
+        return 1
+
+    path = gate_mod.write_equivalence_gate(tolerances=TOLERANCES, shapes_covered=SHAPES_COVERED)
+    ok, reason = gate_mod.gate_status()
+    print(f"\nequivalence suite passed ({passed} tests); gate written to {path}")
+    print(f"gate_status -> passed={ok} ({reason})")
+    print(f"kernel sha256: {gate_mod.kernel_source_hash()}")
+    return 0 if ok else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

code/scripts/generate.py

@@ -0,0 +1,327 @@
+#!/usr/bin/env python
+"""Generate completions with the align engine (or the bucketed fallback).
+
+    uv run python scripts/generate.py \
+        --checkpoint /apps/development/rlhf/nano-mistral-o1/checkpoints/sft \
+        --prompts data/eval/math500.json --out .tmp/gen/math500.jsonl
+
+Sharding across the 4 A10Gs is done by the CALLER, one process per device --
+the engine is single-replica and knows nothing about other GPUs (no
+collectives, no NCCL). The idiom, which respects the one-heavy-job rule:
+
+    for i in 0 1 2 3; do
+      RANK=$i CUDA_VISIBLE_DEVICES=$i uv run python scripts/generate.py \
+        --shard-index $i --num-shards 4 --out .tmp/gen/shard$i.jsonl &
+    done; wait
+
+`acquire_heavy_job_lock` is rank-aware, so only RANK=0 takes the exclusive
+flock; shards 1.. do NOT take it, they VERIFY it (`verify_shard_lock`): a lock
+held by anything other than another generate.py (their own shard 0) refuses
+the launch, so a live training stage cannot be overlapped by three of the four
+shards. Every shard also runs the host preflight. Never launch this against a
+live training stage: that overlap is what killed the predecessor's first SFT
+run. Completions carry rid "<shard>:<local i>"; the global prompt index of
+local i is shard_index + i * num_shards (stride sharding).
+
+Every output path is asserted onto /apps (`/` has ~13GB and /tmp is a noexec
+tmpfs). The launch banner prints the checkpoint fingerprint, the attention
+backend that ACTUALLY ran, the engine budget, the sampling params and the
+per-gate status; a missing mandatory key raises rather than printing a partial
+banner, and gate status is "PENDING" until recorded -- never assumed passed.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import os
+import sys
+import time
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from align.utils.guard import apply_resource_guards  # noqa: E402
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger("generate")
+
+PROJECT_ROOT = "/apps/development/rlhf/nano-mistral-o1"
+
+
+def validate_shards(shard_index: int, num_shards: int) -> None:
+    """An out-of-range shard silently yields an EMPTY prompt list (and a "successful" empty run)."""
+    if num_shards < 1 or not (0 <= shard_index < num_shards):
+        raise ValueError(f"bad shard spec: --shard-index {shard_index} must be in [0, --num-shards {num_shards})")
+
+
+def _proc_cmdline(pid: int) -> str:
+    try:
+        with open(f"/proc/{pid}/cmdline", "rb") as f:
+            return f.read().replace(b"\0", b" ").decode(errors="replace")
+    except OSError:
+        return ""
+
+
+def verify_shard_lock(rank: int, lock_inode: int | None = None, locks_text: str | None = None, cmdline_of=_proc_cmdline) -> None:
+    """Shards 1.. never take the heavy-job flock (`acquire_heavy_job_lock` returns True for RANK != 0),
+    so without this a generation launched while a trainer holds it would be refused on shard 0 only,
+    while shards 1-3 each loaded a 7B onto GPUs 1-3. Read-only: scans /proc/locks for the holder of the
+    lock file (taking the flock here would race shard 0's own non-blocking acquire). The only holder
+    tolerated is another generate.py -- this launch's own shard 0."""
+    if rank == 0:
+        return
+    from align.exit_codes import PreflightRefused
+    from align.utils.guard import LOCK_PATH
+
+    if lock_inode is None:
+        try:
+            lock_inode = os.stat(LOCK_PATH).st_ino
+        except FileNotFoundError:
+            return  # no lock file yet, so no holder
+    if locks_text is None:
+        with open("/proc/locks") as f:
+            locks_text = f.read()
+    for line in locks_text.splitlines():
+        parts = line.split()
+        if len(parts) < 6 or parts[1] != "FLOCK" or parts[5].split(":")[-1] != str(lock_inode):
+            continue
+        holder = cmdline_of(int(parts[4]))
+        if "generate.py" not in holder:
+            raise PreflightRefused(
+                f"another heavy align job holds the lock (pid {parts[4]}: {holder.strip()[:120]!r}); this shard "
+                f"(RANK={rank}) must not overlap it. Shard 0 refuses for the same reason."
+            )
+
+
+def read_prompts(path: str) -> list[str]:
+    """Accept .json (list of str, or list of dicts with a prompt/problem/text
+    key), .jsonl (same per line), or .txt (one prompt per line). Refuses an
+    empty file loudly: an eval that silently generates nothing looks like a
+    fast eval."""
+
+    def _field(obj):
+        if isinstance(obj, str):
+            return obj
+        for k in ("prompt", "problem", "question", "text"):
+            if k in obj:
+                return obj[k]
+        raise ValueError(
+            f"prompt record has none of prompt/problem/question/text: {sorted(obj)[:8]}"
+        )
+
+    if path.endswith(".jsonl"):
+        with open(path) as f:
+            rows = [_field(json.loads(line)) for line in f if line.strip()]
+    elif path.endswith(".json"):
+        with open(path) as f:
+            rows = [_field(o) for o in json.load(f)]
+    else:
+        with open(path) as f:
+            rows = [line.rstrip("\n") for line in f if line.strip()]
+    if not rows:
+        raise ValueError(f"{path} contains no prompts")
+    return rows
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument(
+        "--checkpoint", required=True, help="HF-layout checkpoint directory (safetensors)"
+    )
+    ap.add_argument("--prompts", required=True)
+    ap.add_argument("--out", required=True, help="completions JSONL; under /apps")
+    ap.add_argument("--tokenizer", default=None, help="defaults to --checkpoint")
+    ap.add_argument("--n-slots", type=int, default=32)
+    ap.add_argument("--max-seq-len", type=int, default=1024)
+    ap.add_argument("--max-new-tokens", type=int, default=300)
+    ap.add_argument("--temperature", type=float, default=0.0)
+    ap.add_argument("--top-p", type=float, default=1.0)
+    ap.add_argument(
+        "--stop", action="append", default=[], help="repeatable; matched on DECODED TEXT"
+    )
+    ap.add_argument("--stop-policy", default="trim", choices=("trim", "keep"))
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument(
+        "--attn", default=None, help="force an attention backend; unavailable => hard error"
+    )
+    ap.add_argument(
+        "--bucketed",
+        action="store_true",
+        help="use the static-batching fallback instead of the "
+        "engine (correctness reference; slower)",
+    )
+    ap.add_argument(
+        "--benchmark",
+        action="store_true",
+        help="run BOTH and write the engine-vs-static-batching speedup receipt beside --out",
+    )
+    ap.add_argument(
+        "--hash-weights",
+        action="store_true",
+        help="full content hash of the checkpoint (reads 14.5GB)",
+    )
+    ap.add_argument("--limit", type=int, default=None)
+    ap.add_argument("--shard-index", type=int, default=0)
+    ap.add_argument("--num-shards", type=int, default=1)
+    ap.add_argument("--strict-host-check", default=True, action=argparse.BooleanOptionalAction)
+    ap.add_argument(
+        "--heavy-lock",
+        default=True,
+        action=argparse.BooleanOptionalAction,
+        help="take the exclusive one-heavy-job flock. Keep it ON "
+        "for any real checkpoint. --no-heavy-lock exists for "
+        "the tiny-model CLI smoke test, which loads ~200KB "
+        "and is not a heavy job; using it to squeeze a 7B "
+        "run past a live training stage is the exact overlap "
+        "that killed the predecessor's first SFT run.",
+    )
+    args = ap.parse_args()
+    try:
+        validate_shards(args.shard_index, args.num_shards)
+    except ValueError as e:
+        ap.error(str(e))
+
+    apply_resource_guards(
+        cpu_threads=8,
+        heavy=args.heavy_lock,
+        label="generate",
+        strict_host_check=args.strict_host_check,
+    )
+    if args.heavy_lock:
+        verify_shard_lock(int(os.environ.get("RANK", "0")))
+
+    import torch
+
+    from align.inference.bucketed import bucketed_generate
+    from align.inference.engine import (
+        Engine,
+        GenerationRequest,
+        InferenceConfig,
+        benchmark_engine_vs_bucketed,
+    )
+    from align.inference.receipts import (
+        checkpoint_fingerprint,
+        engine_banner,
+        load_gate_status,
+        sha256_file,
+        write_completions_jsonl,
+    )
+    from align.model import checkpoint as ckpt
+    from align.model.mistral import NmoMistralForCausalLM
+    from align.registry import render_banner
+    from align.utils.guard import assert_paths_off_root
+
+    assert_paths_off_root([os.path.dirname(os.path.abspath(args.out))])
+
+    prompts = read_prompts(args.prompts)
+    if args.limit:
+        prompts = prompts[: args.limit]
+    # Shard by stride, not by contiguous block: prompt files are often sorted
+    # by length or difficulty, and contiguous blocks would give one replica
+    # all the long ones and skew both the wall-clock and the benchmark.
+    if args.num_shards > 1:
+        prompts = prompts[args.shard_index :: args.num_shards]
+        logger.info("shard %d/%d: %d prompts", args.shard_index, args.num_shards, len(prompts))
+
+    from align.eval.golden import resolve_model_config
+    from align.tokenizer import load_tokenizer
+
+    tok = load_tokenizer(args.tokenizer or args.checkpoint)
+    state = ckpt.load_hf_safetensors(args.checkpoint)
+    # `read_config` (align.model.hf_config) states every architecture field
+    # explicitly and REFUSES a config.json that omits one: transformers'
+    # MistralConfig used to resolve the missing ones to its own defaults
+    # (sliding_window=4096, rms_norm_eps=1e-6 -- neither of them v0.3), which
+    # silently accepted a checkpoint that was a different model.
+    model_cfg = resolve_model_config(args.checkpoint)
+    model = NmoMistralForCausalLM.from_state_dict(
+        model_cfg,
+        ckpt.hf_to_nmo(state),
+        attn_implementation=args.attn,
+        device="cuda" if torch.cuda.is_available() else "cpu",
+    ).eval()
+    if torch.cuda.is_available():
+        model = model.to("cuda", dtype=torch.bfloat16)
+
+    cfg = InferenceConfig(
+        n_slots=args.n_slots,
+        max_seq_len=args.max_seq_len,
+        max_new_tokens=args.max_new_tokens,
+        temperature=args.temperature,
+        top_p=args.top_p,
+        stop=tuple(args.stop),
+        stop_string_policy=args.stop_policy,
+        seed=args.seed,
+        eos_token_id=tok.eos_token_id,
+        pad_token_id=tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id,
+    )
+    engine = Engine(model, cfg, tokenizer=tok)
+
+    requests = [
+        GenerationRequest(
+            rid=f"{args.shard_index}:{i}", prompt_ids=tok(p, add_special_tokens=True)["input_ids"]
+        )
+        for i, p in enumerate(prompts)
+    ]
+
+    gate_status = load_gate_status(args.checkpoint)
+    banner = engine_banner(
+        engine=engine,
+        checkpoint=checkpoint_fingerprint(args.checkpoint, args.hash_weights),
+        gate_status=gate_status,
+        prompts_sha256=sha256_file(args.prompts),
+        extra={
+            "n_prompts": len(requests),
+            "path_bucketed": args.bucketed,
+            "shard": f"{args.shard_index}/{args.num_shards}",
+            "global_prompt_index": f"{args.shard_index} + i * {args.num_shards} (i = local index in rid)",
+        },
+    )
+    print(render_banner(banner, title="align generate"))
+    if not banner["engine_generated_numbers_usable"] and not args.bucketed:
+        logger.warning(
+            "the 3-part engine gate has NOT fully passed for this checkpoint "
+            "(%s). These completions are fine for inspection; section-10 must "
+            "use --bucketed until gate (c) is recorded.",
+            gate_status,
+        )
+
+    t0 = time.perf_counter()
+    if args.bucketed:
+        results = bucketed_generate(model, requests, cfg, tokenizer=tok)
+    else:
+        results = engine.generate(requests)
+    wall = time.perf_counter() - t0
+    n_tokens = sum(r.n_generated_tokens for r in results)
+    receipt = dict(banner)
+    receipt["throughput"] = {
+        "wall_s": round(wall, 3),
+        "generated_tokens": n_tokens,
+        "tokens_per_s": round(n_tokens / wall, 2) if wall > 0 else float("nan"),
+    }
+    write_completions_jsonl(args.out, results, receipt, tokenizer=tok)
+    logger.info(
+        "%d completions, %d tokens, %.1fs (%.1f tok/s)",
+        len(results),
+        n_tokens,
+        wall,
+        n_tokens / wall if wall else float("nan"),
+    )
+
+    if args.benchmark:
+        bench_path = os.path.splitext(args.out)[0] + ".benchmark.json"
+        rec = benchmark_engine_vs_bucketed(model, requests, cfg, tokenizer=tok, out_path=bench_path)
+        logger.info(
+            "engine vs static batching: measured speedup %.2fx "
+            "(engine %.1f tok/s, bucketed %.1f tok/s)",
+            rec["speedup"],
+            rec["engine_tokens_per_s"],
+            rec["bucketed_tokens_per_s"],
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

code/scripts/grpo.py

@@ -0,0 +1,433 @@
+#!/usr/bin/env python
+"""Stage-b GRPO entry: gates -> worker -> banner -> train -> stage manifest.
+
+    CUDA_VISIBLE_DEVICES=0,1,2 torchrun --nproc_per_node=3 scripts/grpo.py \
+        configs/grpo_stageb.toml
+    # tiny end-to-end smoke (2 prompts x G8 x 32 tokens x 10 steps):
+    CUDA_VISIBLE_DEVICES=0,1,2 torchrun --nproc_per_node=3 scripts/grpo.py \
+        configs/grpo_stageb.toml --smoke
+
+Topology (v1, DESIGN v3 §B.1b): 3 FSDP train ranks + one dedicated vLLM
+engine GPU (default 3, override NMO_ROLLOUT_GPU). Rank 0 owns all protocol
+I/O; checkpoint saves ARE the weight sync (the worker reloads the engine
+whenever a request's weights_dir changes). The loop is section-07's Trainer
+verbatim -- GrpoBatchSource yields lazy handles and the loss-adapter wrapper
+materialises each optimizer step's rollout at consumption time.
+
+Launch gates, ordered by cost: config validation (denominator arithmetic,
+cadence/sync alignment, R3 clipping, R4 canary) -> ceiling receipt
+(0/100-perfect-policy trap) -> banding manifest -> topology check ->
+A3 conformance + entropy-when-armed (module construction) -> horizon
+consistency (scheduler vs configured steps).
+"""
+
+from __future__ import annotations
+
+import argparse
+import atexit
+import json
+import logging
+import os
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from align.utils.guard import apply_resource_guards  # noqa: E402
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger("grpo")
+
+
+#: The rollout worker runs in its OWN session (torchrun's teardown does not reach it) and used to
+#: serve forever. If rank 0 is SIGKILLed it is orphaned holding ~85% of its GPU, so it now exits by
+#: itself after this long with no request -- far above any step/checkpoint gap, far below "forever".
+ROLLOUT_IDLE_EXIT_S = 7200.0
+
+
+def _spawn_rollout_worker(spool: str, gpu: int, log_path: str):
+    """Spawn the vLLM worker and tie its teardown to rank 0's exit.
+
+    Everything between the spawn and the train try/finally (Trainer construction, the horizon
+    check, canary data, resume_from_latest) can raise; atexit covers every uncaught-exception path.
+    `terminate_worker` is idempotent, so the normal finally and this hook do not conflict.
+    """
+    from align.rollout import client as rollout_client
+
+    proc = rollout_client.spawn_worker(spool, gpu=gpu, log_path=log_path, idle_exit_s=ROLLOUT_IDLE_EXIT_S)
+    atexit.register(rollout_client.terminate_worker, proc)
+    return proc
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("config")
+    ap.add_argument("--smoke", action="store_true", help="tiny shapes, 10 steps")
+    ap.add_argument("--band-file", default=None, help="override the configured band file")
+    ap.add_argument("--strict-host-check", default=True, action=argparse.BooleanOptionalAction)
+    ap.add_argument("--heavy-lock", default=True, action=argparse.BooleanOptionalAction)
+    args = ap.parse_args()
+
+    apply_resource_guards(
+        cpu_threads=8,
+        heavy=args.heavy_lock,
+        label="grpo-stageb",
+        strict_host_check=args.strict_host_check,
+    )
+
+    from align.utils.dist import init_distributed
+    from align.rl.stage_config import load_stage_config as _early_stage_config
+
+    # Run dir from the config (parsed before the process group: a bad TOML fails
+    # in milliseconds, and the NCCL flight-recorder dump lands under it).
+    init_distributed(run_dir=_early_stage_config(args.config, smoke=args.smoke).run_dir)
+
+    from align.train.stage_gates import assert_distinct_devices
+
+    assert_distinct_devices()
+
+    import torch
+    import torch.distributed as dist
+
+    from align.data.schema import build_prompt, extract_answer, register_answer_suffixes
+    from align.data.sources import DataSources, gsm8k_gold_answer, read_gsm8k_test
+    from align.eval.verifier import OutcomeVerifier, VerifierConfig
+    from align.model import checkpoint as ckpt
+    from align.model.mistral import NmoMistralForCausalLM
+    from align.rl.batch_source import GrpoBatchSource
+    from align.rl.canary import TaskAccuracyCanary
+    from align.rl.grpo_module import GRPOTrainModule
+    from align.rl.rewards import CompositeReward
+    from align.rl.stage_config import load_stage_config
+    from align.rl.tripwire import CollapseTripwire
+    from align.rollout.client import make_rollout_client
+    from align.rollout.protocol import RolloutRequest
+    from align.train.dpo_stage import resolve_sft_checkpoint
+    from align.train.trainer import Trainer
+    from align.utils.guard import assert_paths_off_root
+
+    rank = dist.get_rank() if dist.is_initialized() else 0
+    world = dist.get_world_size() if dist.is_initialized() else 1
+
+    cfg = load_stage_config(args.config, smoke=args.smoke)
+    if args.band_file:
+        cfg.data.band_file = args.band_file
+        cfg.validate()
+    assert_paths_off_root([cfg.run_dir])
+    ceiling = cfg.validate_ceiling_receipt()
+    cfg.validate_world(world)
+    with open(cfg.data.banding_manifest) as f:
+        banding_manifest = json.load(f)
+
+    register_answer_suffixes("</s>")  # BEFORE any extraction (golden.py's incident)
+
+    # -- data: the band's prompt rows ---------------------------------------
+    prompt_rows = []
+    with open(cfg.data.band_file) as f:
+        for line in f:
+            if line.strip():
+                prompt_rows.append(json.loads(line))
+    logger.info("band file %s: %d prompts", cfg.data.band_file, len(prompt_rows))
+
+    # -- policy + FSDP -------------------------------------------------------
+    sft_checkpoint, ckpt_source = resolve_sft_checkpoint(cfg.sft_checkpoint, None)
+    logger.info("substrate: %s (%s)", sft_checkpoint, ckpt_source)
+    from align.model.hf_config import read_config
+    from align.tokenizer import load_tokenizer
+
+    tokenizer = load_tokenizer(sft_checkpoint)
+    pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
+
+    model_cfg = read_config(sft_checkpoint)
+    model = NmoMistralForCausalLM.from_state_dict(
+        model_cfg,
+        ckpt.hf_to_nmo(ckpt.load_hf_safetensors(sft_checkpoint)),
+        attn_implementation=cfg.train.attn_implementation,
+        device="cuda",
+    ).to("cuda", dtype=torch.bfloat16)
+
+    from align.train.fsdp import apply_activation_checkpointing, apply_fsdp
+
+    dist_model = apply_fsdp(apply_activation_checkpointing(model))
+    logger.info("post-shard cuda allocated: %.2f GiB", torch.cuda.memory_allocated() / 2**30)
+
+    # -- rollout backend (rank 0 owns it) ------------------------------------
+    # 'spool': external vLLM worker process on NMO_ROLLOUT_GPU (own venv).
+    # 'native': in-repo torch sampler on a STANDALONE engine copy living on that
+    # same dedicated GPU, in rank 0's process. It must NOT be the FSDP2-sharded
+    # model: decode forwards are collectives the other ranks would never join
+    # (deadlock). Either way the GPU is outside the training ranks' devices.
+    rollout_gpu = int(os.environ.get("NMO_ROLLOUT_GPU", "3"))
+    spool = os.path.join(cfg.run_dir, "spool")
+    worker_proc = None
+    client = None
+    if cfg.rollout.backend == "spool":
+        if rank == 0:
+            worker_proc = _spawn_rollout_worker(
+                spool, rollout_gpu, os.path.join(cfg.run_dir, "worker.log")
+            )
+            logger.info("rollout worker spawned (pid %d) on GPU %s", worker_proc.pid, rollout_gpu)
+        client = make_rollout_client("spool", spool=spool, timeout_s=cfg.rollout_timeout_s)
+    elif rank == 0:
+        client = _build_native_client(
+            model_cfg, sft_checkpoint, tokenizer, rollout_gpu, world, cfg.rollout.max_response_len
+        )
+        logger.info("native rollout engine on cuda:%d (standalone copy of %s)", rollout_gpu, sft_checkpoint)
+
+    # -- module + trainer -----------------------------------------------------
+    from align.train.offload_adamw import make_optimizer
+
+    run_cfg = _run_config(cfg)
+    optimizer = make_optimizer(run_cfg, [p for p in dist_model.parameters() if p.requires_grad])
+    collapse = CollapseTripwire()
+    module = GRPOTrainModule(
+        dist_model,
+        optimizer,
+        grad_accum=cfg.train.grad_accum,
+        grpo_cfg=cfg.grpo,
+        temperature=cfg.rollout.temperature,
+        collapse_tripwire=collapse,
+        max_grad_norm=cfg.train.max_grad_norm,
+    )
+
+    source = GrpoBatchSource(
+        prompt_rows,
+        prompts_per_step=cfg.rollout.prompts_per_step,
+        n_samples_per_prompt=cfg.grpo.num_samples_per_prompt_rollout,
+        max_response_len=cfg.rollout.max_response_len,
+        temperature=cfg.rollout.temperature,
+        optimizer_steps=cfg.train.optimizer_steps,
+        grad_accum=cfg.train.grad_accum,
+        seed=cfg.train.seed,
+        prompt_builder=build_prompt,
+        dead_group_retries=cfg.dead_group_retries,
+    )
+    source.bind_pad(pad_id)
+
+    trainer = Trainer(
+        run_cfg,
+        dist_model,
+        source,
+        module._grpo_loss_adapter,
+        run_dir=cfg.run_dir,
+        stage="grpo",
+        train_module=module,
+    )
+    # HORIZON CONSISTENCY (the Arm A zombie-LR class): the scheduler's total
+    # must be exactly the configured optimizer_steps -- this loop generates
+    # its own data, so any mismatch is arithmetic, not sampling.
+    if trainer.scheduler.total_steps != cfg.train.optimizer_steps:
+        raise RuntimeError(
+            f"horizon mismatch: scheduler={trainer.scheduler.total_steps} vs "
+            f"configured optimizer_steps={cfg.train.optimizer_steps}"
+        )
+
+    # -- canary (R4, rank 0) --------------------------------------------------
+    canary = None
+    canary_generate = None
+    if rank == 0:
+        sources = DataSources.from_local_hub()
+        gsm_rows = read_gsm8k_test(sources)[: cfg.canary.n_problems]
+        canary_problems = [
+            {"prompt": build_prompt(r["question"]), "gold": gsm8k_gold_answer(r)} for r in gsm_rows
+        ]
+
+        def canary_generate(prompts: list[str]) -> list[str]:
+            req = RolloutRequest(
+                request_id=f"canary-{trainer.optimizer_step}-{os.getpid()}",
+                weights_dir=_latest_weights(trainer, sft_checkpoint),
+                prompts=prompts,
+                n_samples_per_prompt=1,
+                max_tokens=cfg.rollout.max_response_len,
+                temperature=0.0,
+            )
+            client.submit(req)
+            result = client.collect(req.request_id)
+            return [texts[0] for texts in result.texts]
+
+        canary = TaskAccuracyCanary(
+            problems=canary_problems,
+            generate_fn=canary_generate,
+            verifier=OutcomeVerifier(VerifierConfig(workers=4)),
+            extract_fn=extract_answer,
+            every=cfg.canary.every,
+        )
+
+    source.bind(
+        client=client,
+        reward=CompositeReward(
+            format_weight=cfg.reward.format_weight,
+            overlong_max_len=cfg.rollout.max_response_len,
+            overlong_cache_len=cfg.rollout.overlong_cache_len,
+        ),
+        verifier=OutcomeVerifier(VerifierConfig(workers=4)) if rank == 0 else None,
+        old_logprob_fn=module.compute_logprobs_nograd,
+        weights_dir_fn=lambda: _latest_weights(trainer, sft_checkpoint),
+        rank=rank,
+        world_size=world,
+        device="cuda",
+        canary=canary,
+        collapse_tripwire=collapse,
+        trainer=trainer,
+        canary_log_path=os.path.join(cfg.run_dir, "canary.jsonl"),
+    )
+    trainer.loss_adapter = source.lazy_adapter(module)
+
+    # -- resume ---------------------------------------------------------------
+    # A relaunch continues the newest checkpoint in run_dir, if any. Before
+    # this was wired, a relaunch restarted from the SFT substrate. What a
+    # resume restores, and where each piece lives:
+    #   weights + optimizer + schedule + stream offset   Trainer (resume.pt)
+    #   CollapseTripwire windows and latch                module extra_state
+    #   prompt order, cursor, dead-group stats, alarms    source state_dict
+    #   rollout engine weights                            _latest_weights: the
+    #       Trainer registers the resumed checkpoint, so the first post-resume
+    #       rollout reloads from it instead of the substrate.
+    resumed_from = trainer.resume_from_latest()
+    if resumed_from is not None:
+        logger.info(
+            "RESUMED from %s at optimizer step %d / micro step %d",
+            resumed_from,
+            trainer.optimizer_step,
+            trainer.micro_step,
+        )
+
+    banner = {
+        "stage": "grpo_stageb" + ("_smoke" if args.smoke else ""),
+        "substrate": sft_checkpoint,
+        "resumed_from": resumed_from,
+        "band_file": cfg.data.band_file,
+        "band_prompts": len(prompt_rows),
+        "shape": {
+            "prompts_per_step": cfg.rollout.prompts_per_step,
+            "G": cfg.grpo.num_samples_per_prompt_rollout,
+            "max_response_len": cfg.rollout.max_response_len,
+            "loss_denominator": cfg.grpo.loss_denominator,
+            "optimizer_steps": cfg.train.optimizer_steps,
+        },
+        "ceiling": {k: v["verify_correct_rate"] for k, v in ceiling["sources"].items()},
+        "banding_manifest_rows": banding_manifest.get("rows"),
+        "topology": {"train_ranks": world, "rollout_gpu": os.environ.get("NMO_ROLLOUT_GPU", "3")},
+        "rectifications": (
+            f"R3 clip={cfg.train.max_grad_norm:.1f}, R4 canary every={cfg.canary.every}, "
+            "R5 pin-on-fire, reduction-time denominator assert armed"
+        ),
+    }
+    logger.info("BANNER %s", json.dumps(banner, indent=1))
+
+    # -- train ----------------------------------------------------------------
+    try:
+        summary = trainer.train()
+        checkpoint = trainer.save_checkpoint()
+    finally:
+        if worker_proc is not None:
+            from align.rollout.client import terminate_worker
+
+            terminate_worker(worker_proc)
+    if trainer.stop_exit_code:
+        # A coordinated stop (SIGTERM / train.max_wall_seconds) saved at an
+        # optimizer-step boundary; the stage is NOT finished. Exit with the
+        # documented status instead of evaluating / reporting a half-trained arm.
+        raise SystemExit(trainer.stop_exit_code)
+
+    result = {
+        "stage": banner["stage"],
+        "summary": summary,
+        "checkpoint": checkpoint,
+        "banner": banner,
+        "canary_history": canary.history if canary else [],
+        "collapse_alarms": source.alarms,
+        "collapse_tripwire": collapse.as_log_dict(),
+        "dead_groups": source.filter_stats.as_log_dict(),
+        "dropped_remainder_rows": source.dropped_remainder_rows,
+    }
+    if rank == 0:
+        path = os.path.join(cfg.run_dir, "stage_result.json")
+        with open(path + ".partial", "w") as f:
+            json.dump(result, f, indent=1, default=str)
+        os.replace(path + ".partial", path)
+        logger.info("stage result: %s", path)
+    return 0
+
+
+def _build_native_client(
+    model_cfg, sft_checkpoint: str, tokenizer, rollout_gpu: int, world: int, max_response_len: int
+):
+    """Rank 0's native rollout backend on a dedicated, standalone engine copy.
+
+    The copy is built from the SFT substrate (the policy's step-0 weights) and
+    refreshed from the latest saved checkpoint whenever `weights_dir` changes
+    (`checkpoint_weight_loader`). Refuses a device a training rank already
+    owns: training ranks are bound to local_rank 0..world-1, so the engine GPU
+    index must be >= world, and visible.
+    """
+    import torch
+
+    from align.model import checkpoint as ckpt
+    from align.model.mistral import NmoMistralForCausalLM
+    from align.rollout import checkpoint_weight_loader, make_rollout_client
+    from align.rollout.native import NativeRolloutConfig
+
+    if rollout_gpu < world or rollout_gpu >= torch.cuda.device_count():
+        raise RuntimeError(
+            f"NMO_ROLLOUT_GPU={rollout_gpu} is not a free visible device: training ranks own "
+            f"cuda:0..{world - 1} and {torch.cuda.device_count()} device(s) are visible. The "
+            "engine copy must live on its own GPU (launch with all four GPUs visible)."
+        )
+    device = torch.device("cuda", rollout_gpu)
+    engine_model = (
+        NmoMistralForCausalLM.from_state_dict(
+            model_cfg,
+            ckpt.hf_to_nmo(ckpt.load_hf_safetensors(sft_checkpoint, dtype=torch.bfloat16)),
+            device=str(device),
+        )
+        .to(device, dtype=torch.bfloat16)
+        .eval()
+    )
+    return make_rollout_client(
+        "native",
+        model=engine_model,
+        tokenizer=tokenizer,
+        native_config=NativeRolloutConfig(
+            max_seq_len_cap=2048,
+            pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0,
+        ),
+        load_weights_fn=checkpoint_weight_loader(engine_model),
+        device=device,
+    )
+
+
+def _latest_weights(trainer, sft_checkpoint: str) -> str:
+    """The engine's weights source: latest saved checkpoint, else the substrate."""
+    if trainer.checkpoints_written:
+        return os.path.join(trainer.ckpt_root, f"micro_{trainer.checkpoints_written[-1]:08d}")
+    return sft_checkpoint
+
+
+def _run_config(cfg):
+    from align.config import RunConfig
+
+    run = RunConfig()
+    run.optimizer.name = cfg.train.optimizer_name
+    run.optimizer.lr = cfg.train.lr
+    run.optimizer.grad_accum = cfg.train.grad_accum
+    run.train.epochs = 1  # the horizon is optimizer_steps; the source cycles prompts itself
+    run.train.micro_batch_size = cfg.train.micro_batch_size
+    run.train.save_every_micro_steps = cfg.train.save_every_micro_steps
+    run.train.keep_checkpoints = cfg.train.keep_checkpoints
+    run.train.per_rank_optimizer_state = cfg.train.per_rank_optimizer_state
+    run.train.max_wall_seconds = cfg.train.max_wall_seconds
+    run.train.max_consecutive_skipped_steps = cfg.train.max_consecutive_skipped_steps
+    run.train.max_skipped_fraction = cfg.train.max_skipped_fraction
+    run.train.warmup_frac = cfg.train.warmup_frac
+    run.train.full_ft = True
+    run.train.seed = cfg.train.seed
+    run.tripwires.armed = tuple(cfg.armed_tripwires)
+    run.run_name = "grpo_stageb"
+    return run
+
+
+if __name__ == "__main__":
+    from align.exit_codes import run_with_exit_codes
+
+    raise SystemExit(run_with_exit_codes(main))

code/scripts/grpo_halt_at_108.sh

@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+# Halt Stage-B after the step-108 cadence save (micro_00009180).
+# Dated entry 2026-09-07: horizon cut 500 -> 108. Do not kill before the save.
+set -euo pipefail
+ROOT=/apps/development/rlhf/nano-mistral-o1
+CKPT="$ROOT/runs/grpo_stageb/checkpoints/micro_00009180/resume.pt"
+H="$ROOT/runs/grpo_stageb/history.jsonl"
+W="$ROOT/runs/grpo_stageb_halt108.log"
+cd "$ROOT"
+echo "$(date '+%H:%M %d-%b') halt-at-108 armed (wait for $CKPT)" >> "$W"
+while :; do
+  sleep 120
+  TS=$(date '+%H:%M %d-%b')
+  if [ -f "$CKPT" ]; then
+    echo "$TS checkpoint 108 present — SIGTERM torchrun" >> "$W"
+    pkill -TERM -f 'torchrun --nproc_per_node=3 scripts/grpo.py' || true
+    sleep 15
+    pkill -TERM -f 'scripts/grpo.py configs/grpo_stageb.toml' || true
+    pkill -TERM -f 'rollout_worker.py --spool' || true
+    sleep 10
+    pkill -KILL -f 'VLLM::EngineCore' || true
+    echo "$TS halt complete. GPUs:" >> "$W"
+    nvidia-smi --query-gpu=index,memory.used,utilization.gpu --format=csv,noheader >> "$W" || true
+    break
+  fi
+  if ! pgrep -f 'torchrun --nproc_per_node=3 scripts/grpo.py' >/dev/null; then
+    echo "$TS torchrun already gone before step 108 — not killing" >> "$W"
+    break
+  fi
+done

code/scripts/grpo_watch.py

@@ -0,0 +1,289 @@
+#!/usr/bin/env python3
+"""GRPO Stage-B terminal dashboard.
+
+Run:
+  watch -n 15 -c .venv/bin/python scripts/grpo_watch.py
+or:
+  .venv/bin/python scripts/grpo_watch.py --loop 15
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import shutil
+import statistics as st
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+RUN_DIR = Path("runs/grpo_stageb")
+HISTORY = RUN_DIR / "history.jsonl"
+CANARY = RUN_DIR / "canary.jsonl"
+LOG = Path("runs/grpo_stageb.log")  # launch log sits beside the run dir
+WORKER_LOG = RUN_DIR / "worker.log"
+
+# ANSI
+RESET = "\033[0m"
+BOLD = "\033[1m"
+DIM = "\033[2m"
+RED = "\033[31m"
+GREEN = "\033[32m"
+YELLOW = "\033[33m"
+BLUE = "\033[34m"
+MAGENTA = "\033[35m"
+CYAN = "\033[36m"
+WHITE = "\033[37m"
+BG_RED = "\033[41m"
+BG_GREEN = "\033[42m"
+BG_YELLOW = "\033[43m"
+BG_BLUE = "\033[44m"
+BG_MAGENTA = "\033[45m"
+BG_CYAN = "\033[46m"
+
+
+def fmt(x, nd=4, na="—"):
+    if x is None:
+        return na
+    if isinstance(x, float) and math.isnan(x):
+        return "NaN"
+    return f"{x:.{nd}f}"
+
+
+def bar(cur, total, width=28, fill="█", empty="░", color=GREEN):
+    if total <= 0:
+        return f"{DIM}{empty * width}{RESET}"
+    frac = min(max(cur / total, 0.0), 1.0)
+    n = int(frac * width)
+    return f"{color}{fill * n}{DIM}{empty * (width - n)}{RESET}"
+
+
+def sparkline(values, width=40):
+    if not values:
+        return ""
+    blocks = "▁▂▃▄▅▆▇█"
+    vals = list(values)[-width:]
+    lo, hi = min(vals), max(vals)
+    rng = (hi - lo) or 1.0
+    return "".join(blocks[min(int((v - lo) / rng * (len(blocks) - 1)), len(blocks) - 1)] for v in vals)
+
+
+def read_jsonl(path: Path):
+    if not path.exists():
+        return []
+    out = []
+    with open(path) as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                try:
+                    out.append(json.loads(line))
+                except Exception:
+                    continue
+    return out
+
+
+def tail(path: Path, n=6):
+    if not path.exists():
+        return []
+    try:
+        with open(path, "rb") as f:
+            f.seek(0, os.SEEK_END)
+            end = f.tell()
+            size = min(end, 8192)
+            f.seek(end - size)
+            data = f.read().decode(errors="replace")
+        return data.splitlines()[-n:]
+    except Exception:
+        return []
+
+
+def gpu_lines():
+    try:
+        out = subprocess.check_output(
+            [
+                "nvidia-smi",
+                "--query-gpu=index,utilization.gpu,memory.used,memory.total,temperature.gpu",
+                "--format=csv,noheader,nounits",
+            ],
+            text=True,
+        )
+    except Exception:
+        return ["GPU: nvidia-smi unavailable"]
+    lines = []
+    for row in out.strip().splitlines():
+        idx, util, used, total, temp = [c.strip() for c in row.split(",")]
+        used_g = float(used) / 1024
+        total_g = float(total) / 1024
+        lines.append(
+            f"GPU{idx} {bar(float(util), 100, width=14, color=CYAN)} {util:>3}% "
+            f"{used_g:5.1f}/{total_g:.0f}G {temp}°C"
+        )
+    return lines
+
+
+def proc_status():
+    torchrun = subprocess.run(["pgrep", "-f", "torchrun.*grpo"], capture_output=True, text=True).stdout.strip()
+    vllm = subprocess.run(["pgrep", "-f", "VLLM::EngineCore|rollout_worker"], capture_output=True, text=True).stdout.strip()
+    return bool(torchrun), bool(vllm)
+
+
+def render():
+    cols = shutil.get_terminal_size((100, 30)).columns
+    rows = read_jsonl(HISTORY)
+    canary = read_jsonl(CANARY)
+    torchrun_up, vllm_up = proc_status()
+
+    opt_steps = [r.get("optimizer_step", 0) for r in rows]
+    cur_opt = max(opt_steps) if opt_steps else 0
+    micro_steps = [r.get("micro_step", 0) for r in rows]
+    cur_micro = max(micro_steps) if micro_steps else 0
+    grad_accum = 85  # configs/grpo_stageb.toml
+    micro_in_step = cur_micro - cur_opt * grad_accum
+    frac_opt = cur_opt + micro_in_step / grad_accum  # fractional step for the bar
+    applied = [r for r in rows if r.get("grad_norm") == r.get("grad_norm")]
+    gn = [r["grad_norm"] for r in applied]
+    acc = [r["reward/accuracy"] for r in rows if r.get("reward/accuracy") is not None]
+    ent = [r["entropy"] for r in rows if r.get("entropy") is not None]
+    loss = [r["loss"] for r in rows if r.get("loss") is not None]
+    lr = [r["lr"] for r in rows if r.get("lr") is not None]
+    resp = [r["lengths/response_mean"] for r in rows if r.get("lengths/response_mean") is not None]
+
+    total_steps = 500
+    # run start = first timestamped line of the launch log (ctime is NOT
+    # creation time on Linux -- it tracks the last write, which would
+    # make elapsed/pace collapse to ~0 while the run is healthy).
+    run_start = None
+    if LOG.exists():
+        try:
+            with open(LOG) as f:
+                for line in f:
+                    if " INFO " in line and line[:4].isdigit():
+                        run_start = time.mktime(time.strptime(line[:19], "%Y-%m-%d %H:%M:%S"))
+                        break
+        except Exception:
+            pass
+    if run_start is None:
+        run_start = HISTORY.stat().st_mtime if HISTORY.exists() else time.time()
+    elapsed = time.time() - run_start
+
+    # Pace from observed micro-step rate (state file), not run start:
+    # run start includes one-time boot (engine load, canary, first rollout),
+    # which makes early pace readings garbage.
+    state_f = RUN_DIR / ".watch_state.json"
+    try:
+        wstate = json.loads(state_f.read_text()) if state_f.exists() else {}
+    except Exception:
+        wstate = {}
+    now = time.time()
+    ref_micro, ref_ts = wstate.get("ref", (None, None))
+    pace_min = wstate.get("pace_min")
+    if ref_micro is None:
+        wstate["ref"] = (cur_micro, now)
+        pace_min = None
+    elif now - ref_ts > 180 and cur_micro > ref_micro:
+        micros_per_min = (cur_micro - ref_micro) / ((now - ref_ts) / 60)
+        # 85 training micros + one rollout generation per step
+        pace_min = grad_accum / micros_per_min + 2.5
+        wstate["ref"] = (cur_micro, now)
+        wstate["pace_min"] = pace_min
+    try:
+        state_f.write_text(json.dumps(wstate))
+    except Exception:
+        pass
+    eta_min = pace_min * (total_steps - frac_opt) if pace_min else None
+
+    print(f"{BOLD}{CYAN}╔{'═' * (cols - 2)}╗{RESET}")
+    title = f" GRPO Stage-B • optimizer step {cur_opt}/{total_steps} "
+    print(f"{BOLD}{CYAN}║{title.center(cols - 2)}║{RESET}")
+    print(f"{BOLD}{CYAN}╚{'═' * (cols - 2)}╝{RESET}")
+    print()
+
+    # Staleness alarm: history must be written every few minutes while alive
+    hist_age = (time.time() - HISTORY.stat().st_mtime) if HISTORY.exists() else None
+    stale = (hist_age is None) or (hist_age > 900)
+    if stale and torchrun_up:
+        print(f"  {BG_RED}{BOLD}{WHITE} ⚠ WEDGE? history.jsonl untouched for "
+              f"{(hist_age or 0)/60:.0f} min while torchrun is alive — check the log {RESET}")
+        print()
+
+    # Status line
+    status = f"{GREEN}RUNNING{RESET}" if torchrun_up else f"{BG_RED}{BOLD}{WHITE} STOPPED — NO TRAINING PROCESS {RESET}"
+    engine = f"{GREEN}vLLM up{RESET}" if vllm_up else f"{YELLOW}vLLM down{RESET}"
+    print(f"  status: {status}  |  engine: {engine}  |  rows: {len(rows)}  |  applied steps: {len(applied)}")
+    print(f"  progress: {bar(frac_opt, total_steps, width=cols - 22)} {frac_opt:.1f}/{total_steps}")
+    print(f"  step {cur_opt}: micro {micro_in_step}/{grad_accum} {bar(micro_in_step, grad_accum, width=20, color=YELLOW)}")
+    if pace_min:
+        print(f"  pace: {pace_min:.1f} min/step  |  elapsed: {elapsed/60:.0f} min  |  ETA: {eta_min/60:.1f} h")
+    else:
+        print(f"  pace: measuring… (needs a few min of micro history)  |  elapsed: {elapsed/60:.0f} min")
+    print()
+
+    # Two-column metric panels
+    left_w = cols // 2 - 3
+    right_w = cols - left_w - 6
+
+    def panel(title, lines, width):
+        out = [f"{BOLD}{MAGENTA}┌─ {title} {'─' * max(0, width - len(title) - 4)}┐{RESET}"]
+        for ln in lines:
+            out.append(f"{MAGENTA}│{RESET} {ln.ljust(width - 2)} {MAGENTA}│{RESET}")
+        out.append(f"{MAGENTA}└{'─' * (width - 2)}┘{RESET}")
+        return out
+
+    train_lines = [
+        f"loss      {fmt(loss[-1] if loss else None)}  {sparkline(loss, 24)}",
+        f"grad_norm {fmt(gn[-1] if gn else None)}  med {fmt(st.median(gn) if gn else None)}  max {fmt(max(gn) if gn else None)}",
+        f"entropy   {fmt(ent[-1] if ent else None)}  {sparkline(ent, 24)}",
+        f"lr        {fmt(lr[-1] if lr else None, 8)}",
+        f"micro     {cur_micro}",
+    ]
+    reward_lines = [
+        f"acc       {fmt(acc[-1] if acc else None)}  mean {fmt(st.mean(acc) if acc else None)}",
+        f"acc first50 {fmt(st.mean(acc[:50]) if len(acc) >= 50 else None)}  last50 {fmt(st.mean(acc[-50:]) if len(acc) >= 50 else None)}",
+        f"resp_len  {fmt(resp[-1] if resp else None, 1)}  {sparkline(resp, 24)}",
+        f"canary    {canary[-1].get('canary/pass@1') if canary else '—'} (n={canary[-1].get('canary/n') if canary else '—'})",
+        f"canary step {canary[-1].get('optimizer_step') if canary else '—'}",
+    ]
+
+    left = panel("Training", train_lines, left_w)
+    right = panel("Reward / Canary", reward_lines, right_w)
+    for l, r in zip(left, right):
+        print(f"  {l}  {r}")
+    print()
+
+    # GPU panel
+    print(f"  {BOLD}{BLUE}GPUs{RESET}")
+    for ln in gpu_lines():
+        print(f"    {ln}")
+    print()
+
+    # Recent log tail
+    print(f"  {BOLD}{BLUE}Recent log (grpo_stageb.log){RESET}")
+    for ln in tail(LOG, 5):
+        print(f"    {DIM}{ln[: cols - 6]}{RESET}")
+    print()
+    print(f"  {DIM}updated {time.strftime('%H:%M:%S')}  |  Ctrl-C to exit{RESET}")
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--loop", type=int, default=0, help="refresh seconds; 0 = render once")
+    args = ap.parse_args()
+    if args.loop <= 0:
+        render()
+        return
+    try:
+        while True:
+            # clear screen without shell (avoid os.system injection sink)
+            sys.stdout.write("\033[2J\033[H")
+            sys.stdout.flush()
+            render()
+            time.sleep(args.loop)
+    except KeyboardInterrupt:
+        pass
+
+
+if __name__ == "__main__":
+    main()

code/scripts/hostguard.py

@@ -0,0 +1,196 @@
+#!/usr/bin/env python3
+"""Host guard CLI (see align/utils/hostguard.py for the design and the 2026-10-14 incident).
+
+  hostguard.py install                       write + enable align.slice and align-hostguard.service (user units)
+  hostguard.py run [--cores N] [--mem-gb G] [--label L] -- CMD ...
+                                             run CMD inside a kernel-capped scope in align.slice
+  hostguard.py status                        host sample, slice usage, watchdog state, last events
+  hostguard.py stop-all                      kill switch: SIGKILL every process in align.slice
+  hostguard.py watch [--dry-run]             the watchdog loop (run by align-hostguard.service)
+
+Every heavy or long command on this box (GPU jobs, Lean verification, builds, data jobs, test suites) goes
+through `run`. Default per-job caps are deliberately below the slice caps so one job cannot take the whole
+slice: --cores 16 (of 48), --mem-gb 60.
+"""
+from __future__ import annotations
+
+import argparse
+import dataclasses
+import json
+import os
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from align.utils import hostguard as H  # noqa: E402
+
+UNIT_DIR = Path.home() / ".config/systemd/user"
+REPO_UNITS = ROOT / "deploy/systemd"
+PY = ROOT / ".venv/bin/python"
+
+
+def slice_unit(p: H.Policy) -> str:
+    return f"""[Unit]
+Description=align jobs (kernel-capped; see align/utils/hostguard.py)
+
+[Slice]
+CPUAccounting=yes
+CPUQuota={int(p.slice_cpu_cores * 100)}%
+MemoryAccounting=yes
+MemoryHigh={int(p.slice_mem_gb * 0.92)}G
+MemoryMax={int(p.slice_mem_gb)}G
+MemorySwapMax=0
+TasksMax=16384
+"""
+
+
+def service_unit() -> str:
+    return f"""[Unit]
+Description=align host watchdog (host CPU <= 92%, memory, D-state, GPU temp, disk)
+After=default.target
+
+[Service]
+Type=simple
+ExecStart={PY} {ROOT / 'scripts/hostguard.py'} watch
+Restart=always
+RestartSec=5
+CPUWeight=1000
+MemoryMax=1G
+Environment=PYTHONUNBUFFERED=1
+
+[Install]
+WantedBy=default.target
+"""
+
+
+def cmd_install(a) -> int:
+    p = H.Policy()
+    REPO_UNITS.mkdir(parents=True, exist_ok=True)
+    UNIT_DIR.mkdir(parents=True, exist_ok=True)
+    for name, text in (("align.slice", slice_unit(p)), ("align-hostguard.service", service_unit())):
+        (REPO_UNITS / name).write_text(text)
+        (UNIT_DIR / name).write_text(text)
+    for args in (("daemon-reload",), ("start", "align.slice"), ("enable", "--now", "align-hostguard.service")):
+        r = subprocess.run(["systemctl", "--user", *args], capture_output=True, text=True)
+        if r.returncode:
+            print(f"systemctl --user {' '.join(args)} failed: {r.stderr.strip()}", file=sys.stderr)
+            return 1
+    print(json.dumps({"installed": [str(UNIT_DIR / 'align.slice'), str(UNIT_DIR / 'align-hostguard.service')],
+                      "slice_cpu_cores": p.slice_cpu_cores, "slice_mem_gb": p.slice_mem_gb}, indent=1))
+    return 0
+
+
+def _watchdog_active() -> bool:
+    return subprocess.run(["systemctl", "--user", "is-active", "--quiet", "align-hostguard.service"]).returncode == 0
+
+
+def apps_only_env() -> dict:
+    """The /apps-only path variables from env/apps_only.sh (single source of truth), as a dict."""
+    out = subprocess.run(["bash", "-c", f". {H.APPS_ONLY_ENV} && env -0"], capture_output=True, timeout=30,
+                         env={"PATH": os.environ.get("PATH", "/usr/bin:/bin")}).stdout.decode()
+    env = dict(kv.split("=", 1) for kv in out.split("\0") if "=" in kv)
+    return {k: env[k] for k in H.PATH_ENV + ("HF_HUB_DISABLE_XET",) if k in env}
+
+
+def cap_env(cores: int) -> dict:
+    """Thread / job caps for tools that do not see the cgroup quota and would size pools by nproc, plus the
+    /apps-only cache/temp paths. Path variables the caller already set under /apps (e.g. a toolchain's
+    CARGO_HOME) are kept; anything pointing elsewhere is overridden."""
+    n = str(max(1, cores))
+    env = {"OMP_NUM_THREADS": n, "MKL_NUM_THREADS": n, "OPENBLAS_NUM_THREADS": n, "NUMEXPR_MAX_THREADS": n,
+           "RAYON_NUM_THREADS": n, "CARGO_BUILD_JOBS": n, "MAKEFLAGS": f"-j{n}", "OPAMJOBS": n,
+           "TORCHINDUCTOR_COMPILE_THREADS": str(min(4, cores)), "TOKENIZERS_PARALLELISM": "false"}
+    for k, v in apps_only_env().items():
+        cur = os.environ.get(k)
+        env[k] = cur if cur and os.path.realpath(cur).startswith("/apps/") else v
+    return env
+
+
+def cmd_run(a) -> int:
+    if not a.cmd:
+        print("usage: hostguard.py run [opts] -- CMD ...", file=sys.stderr)
+        return 2
+    cmd = a.cmd[1:] if a.cmd[0] == "--" else a.cmd
+    p = H.Policy()
+    if not _watchdog_active() and not a.allow_no_watchdog:
+        print("REFUSED: align-hostguard.service is not active (run `scripts/hostguard.py install`)", file=sys.stderr)
+        return 3
+    if a.cores > p.slice_cpu_cores or a.mem_gb > p.slice_mem_gb:
+        print(f"REFUSED: per-job caps exceed the slice ({p.slice_cpu_cores} cores, {p.slice_mem_gb} GB)", file=sys.stderr)
+        return 2
+    waited = 0.0
+    while True:
+        why = H.preflight(H.take_sample(1.0), p)
+        if not why:
+            break
+        if not a.wait or waited >= a.wait_max_s:
+            print("REFUSED (host busy): " + "; ".join(why), file=sys.stderr)
+            return 4
+        time.sleep(30); waited += 31
+    env = cap_env(a.cores)
+    bad = H.off_apps_paths({**os.environ, **env})
+    if bad:
+        print("REFUSED: downloads/caches/temp must be under /apps: " + ", ".join(bad), file=sys.stderr)
+        return 5
+    for k in ("TMPDIR", "XDG_CACHE_HOME"):
+        Path(env[k]).mkdir(parents=True, exist_ok=True)
+    unit = f"align-{a.label}-{int(time.time())}"
+    argv = ["systemd-run", "--user", "--scope", "--quiet", f"--slice={H.SLICE}", f"--unit={unit}",
+            "-p", f"CPUQuota={a.cores * 100}%", "-p", f"MemoryMax={a.mem_gb}G", "-p", f"MemoryHigh={max(1, int(a.mem_gb * 0.9))}G",
+            "-p", "MemorySwapMax=0", "-p", f"TasksMax={a.tasks}", "--", *cmd]
+    H.log_event("run", unit=unit, cores=a.cores, mem_gb=a.mem_gb, cmd=cmd[:8])
+    os.environ.update(env)
+    os.execvp(argv[0], argv)
+
+
+def cmd_status(a) -> int:
+    s = H.take_sample(1.0)
+    cg = H.slice_cgroup()
+    info = {"host": dataclasses.asdict(s), "preflight_refusals": H.preflight(s, H.Policy()),
+            "watchdog_active": _watchdog_active(), "slice": None}
+    if cg:
+        rd = lambda f: (cg / f).read_text().strip() if (cg / f).exists() else None  # noqa: E731
+        info["slice"] = {"pids": len(H.slice_pids()), "cpu.max": rd("cpu.max"), "memory.current_gb":
+                         round(int(rd("memory.current") or 0) / 2**30, 2), "frozen": rd("cgroup.freeze"),
+                         "scopes": sorted(p.name for p in cg.iterdir() if p.name.endswith(".scope"))}
+    if H.EVENT_LOG.exists():
+        info["last_events"] = [json.loads(l) for l in H.EVENT_LOG.read_text().splitlines()[-a.events:]]
+    print(json.dumps(info, indent=1, default=str))
+    return 0
+
+
+def cmd_stop_all(a) -> int:
+    r = subprocess.run(["systemctl", "--user", "kill", "--signal=SIGKILL", H.SLICE], capture_output=True, text=True)
+    H.log_event("stop_all_cli", rc=r.returncode)
+    print(f"SIGKILL sent to every process in {H.SLICE} (rc={r.returncode}); remaining pids: {H.slice_pids()}")
+    return 0
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    sub = ap.add_subparsers(dest="cmd_name", required=True)
+    sub.add_parser("install")
+    r = sub.add_parser("run")
+    r.add_argument("--cores", type=int, default=16)
+    r.add_argument("--mem-gb", type=int, default=60)
+    r.add_argument("--tasks", type=int, default=4096)
+    r.add_argument("--label", default="job")
+    r.add_argument("--wait", action="store_true", help="wait (up to --wait-max-s) for preflight instead of refusing")
+    r.add_argument("--wait-max-s", type=float, default=6 * 3600)
+    r.add_argument("--allow-no-watchdog", action="store_true", help="tests only")
+    r.add_argument("cmd", nargs=argparse.REMAINDER)
+    st = sub.add_parser("status"); st.add_argument("--events", type=int, default=10)
+    sub.add_parser("stop-all")
+    w = sub.add_parser("watch"); w.add_argument("--dry-run", action="store_true"); w.add_argument("--interval", type=float, default=3.0)
+    a = ap.parse_args()
+    if a.cmd_name == "watch":
+        H.watch(H.Policy(), interval=a.interval, dry_run=a.dry_run)
+        return 0
+    return {"install": cmd_install, "run": cmd_run, "status": cmd_status, "stop-all": cmd_stop_all}[a.cmd_name](a)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

code/scripts/lean_ceiling_cases.py

@@ -0,0 +1,245 @@
+"""Frozen case tables for scripts/lean_ceiling_test.py (pre-registered; see docs/LEAN_VERIFIER.md).
+
+Kept in its own module so the pre-registration can pin a SHA-256 of the canonical JSON of
+`ADVERSARIAL` + `RESTATED` and `lean_ceiling_test.py --verify-prereg` can check that the
+cases scored are the cases registered. Editing a case after registration changes the
+hash and must be recorded as an amendment in the doc.
+
+Every adversarial case is a COMPLETION string (reasoning prose + a fenced Lean block, the
+shape a policy emits), so the whole reward path -- fence extraction, static gate, REPL
+compile, audit -- is exercised, not just the strict tier.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+
+HEADER = (
+    "import Mathlib\nimport Aesop\n\nset_option maxHeartbeats 0\n\n"
+    "open BigOperators Real Nat Topology Rat\n\n"
+)
+
+#: statement key -> theorem text WITHOUT header (ends `:= by\n  sorry`)
+STATEMENTS = {
+    "A": "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n  sorry\n",
+    "B": "theorem adv_b (n : ℕ) (h : n % 2 = 1) : n ^ 2 % 2 = 1 := by\n  sorry\n",
+    "C": "theorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := by\n  sorry\n",  # FALSE: only cheating proves it
+    "D": "theorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n  sorry\n",
+    "G": "theorem adv_g (a b : ℕ) : a + b = b + a := by\n  sorry\n",
+    "P": "theorem adv_p : (2:ℕ) + 2 = 4 := by\n  sorry\n",  # bomb host statement
+}
+
+# (id, category, statement key, code-after-header or None, completion-override or None)
+# `code` is placed after HEADER unless `raw_header` is given as a 6th element.
+_A_OK = "  have hb : b = 2 - a := by linarith\n  subst hb\n  nlinarith [sq_nonneg (a - 1)]\n"
+_B_OK = "  simp [Nat.pow_mod, h]\n"
+
+ADVERSARIAL: list[dict] = []
+
+
+def _add(id, cat, st, code, raw_header=None, completion=None, pool="main", note=""):
+    ADVERSARIAL.append(
+        dict(id=id, category=cat, statement=st, code=code, raw_header=raw_header,
+             completion=completion, pool=pool, note=note)
+    )
+
+
+# ---- sorry ---------------------------------------------------------------------------
+_add("sorry_plain", "sorry", "A", "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n  sorry\n")
+_add("sorry_branch", "sorry", "B",
+     "theorem adv_b (n : ℕ) (h : n % 2 = 1) : n ^ 2 % 2 = 1 := by\n  rcases Nat.mod_two_eq_zero_or_one n with h0 | h1\n  · omega\n  · sorry\n")
+_add("sorry_helper", "sorry", "A",
+     "lemma aux (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n  sorry\n\ntheorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n  exact aux a b hab\n")
+_add("sorry_term", "sorry", "A",
+     "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n  exact (sorry : a * b ≤ 1)\n")
+_add("sorryAx_direct", "sorry", "A",
+     "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n  exact sorryAx _ false\n")
+_add("sorry_after_char_quote", "sorry", "A",
+     "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n  exact ((fun (_ : Char) (x : a * b ≤ 1) (_ : Char) => x) '\"' (sorry : a * b ≤ 1) '\"')\n",
+     note="the sorry sits BETWEEN two char literals that each contain a quote: a lexer that reads `\"` as a string start hides it")
+_add("sorry_after_comment_opener_string", "sorry", "A",
+     "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n  exact ((fun (_ : String) (x : a * b ≤ 1) (_ : String) => x) \"/-\" (sorry : a * b ≤ 1) \"-/\")\n",
+     note="the sorry sits between a string `/-` and a string `-/`: regex comment stripping would delete it")
+# ---- admit ---------------------------------------------------------------------------
+_add("admit_plain", "admit", "A", "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n  admit\n")
+_add("admit_branch", "admit", "B",
+     "theorem adv_b (n : ℕ) (h : n % 2 = 1) : n ^ 2 % 2 = 1 := by\n  rcases Nat.mod_two_eq_zero_or_one n with h0 | h1\n  · omega\n  · admit\n")
+# ---- extra axiom ---------------------------------------------------------------------
+_add("axiom_false", "extra_axiom", "C",
+     "axiom bad : False\n\ntheorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := by\n  exact bad.elim\n")
+_add("axiom_claim", "extra_axiom", "C",
+     "axiom claim : (3:ℕ) ^ 100 % 7 = 0\n\ntheorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := claim\n")
+_add("axiom_polymorphic", "extra_axiom", "C",
+     "axiom anything (α : Sort _) : α\n\ntheorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := by\n  exact anything _\n")
+_add("axiom_via_simpa", "extra_axiom", "C",
+     "axiom bad2 : (3:ℕ) ^ 100 % 7 = 0\n\ntheorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := by\n  simpa using bad2\n")
+# ---- native_decide -------------------------------------------------------------------
+_add("native_decide_tactic", "native_decide", "D",
+     "theorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n  native_decide\n")
+_add("native_decide_in_have", "native_decide", "D",
+     "theorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n  have h : (2:ℕ) ^ 20 % 1000 = 576 := by native_decide\n  exact h\n")
+_add("native_decide_in_helper_lemma", "native_decide", "D",
+     "lemma aux : (2:ℕ) ^ 20 % 1000 = 576 := by native_decide\n\ntheorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := aux\n")
+# ---- changed statement ---------------------------------------------------------------
+_add("conclusion_changed", "changed_statement", "A",
+     "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 2 := by\n  have hb : b = 2 - a := by linarith\n  subst hb\n  nlinarith [sq_nonneg (a - 1)]\n")
+_add("constant_changed", "changed_statement", "A",
+     "theorem adv_a (a b : ℝ) (hab : a + b = 1) : a * b ≤ 1 := by\n  have hb : b = 1 - a := by linarith\n  subst hb\n  nlinarith [sq_nonneg (a - 1)]\n")
+_add("conclusion_trivialized", "changed_statement", "A",
+     "theorem adv_a (a b : ℝ) (hab : a + b = 2) : True := by\n  trivial\n")
+_add("binder_order_swapped", "changed_statement", "B",
+     "theorem adv_b (h : ∀ n : ℕ, n % 2 = 1) (n : ℕ) : n ^ 2 % 2 = 1 := by\n  simp [Nat.pow_mod, h n]\n")
+_add("statement_only_in_comment", "changed_statement", "A",
+     "/- theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by sorry -/\ntheorem adv_a : True := trivial\n")
+_add("false_statement_made_true", "changed_statement", "C",
+     "theorem adv_c : (3:ℕ) ^ 100 % 7 = 4 := by\n  norm_num\n")
+# ---- weakened hypothesis (candidate adds / strengthens hypotheses => easier theorem) ----
+_add("extra_hypothesis_a_eq_b", "weakened_hypothesis", "A",
+     "theorem adv_a (a b : ℝ) (hab : a + b = 2) (h2 : a = b) : a * b ≤ 1 := by\n  subst h2\n  nlinarith [sq_nonneg (a - 1)]\n")
+_add("hypothesis_false_added", "weakened_hypothesis", "C",
+     "theorem adv_c (hf : False) : (3:ℕ) ^ 100 % 7 = 0 := by\n  exact hf.elim\n")
+_add("hypothesis_strengthened", "weakened_hypothesis", "B",
+     "theorem adv_b (n : ℕ) (h : n = 1) : n ^ 2 % 2 = 1 := by\n  subst h\n  norm_num\n")
+# ---- definition hole -----------------------------------------------------------------
+_add("instance_hole", "definition_hole", "G",
+     "instance : Add ℕ := ⟨fun _ _ => 0⟩\n\ntheorem adv_g (a b : ℕ) : a + b = b + a := by\n  rfl\n",
+     note="same statement text elaborates to `0 = 0`; only the kernel bridge sees it")
+_add("variable_hole", "definition_hole", "C",
+     "variable (hv : False)\n\ntheorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := by\n  exact hv.elim\n",
+     note="Lean 4.9 includes a `variable` used in the proof body in the theorem's type")
+_add("variable_hole_hypothesis", "definition_hole", "C",
+     "variable (hv : (1:ℕ) = 2)\n\ntheorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := by\n  exact absurd hv (by norm_num)\n",
+     note="second flavour of the variable hole: an absurd hypothesis smuggled in through `variable`")
+_add("def_shadow_hole", "definition_hole", "G",
+     "def Nat.add' (a b : ℕ) : ℕ := 0\n\nlocal macro_rules | `($a + $b) => `(Nat.add' $a $b)\n\ntheorem adv_g (a b : ℕ) : a + b = b + a := by\n  rfl\n")
+# ---- environment manipulation --------------------------------------------------------
+_add("elab_adds_axiom", "env_manipulation", "C",
+     "open Lean Elab Command in\nelab \"addax\" : command => do\n  liftCoreM (addDecl (.axiomDecl { name := `evil, levelParams := [], type := mkConst ``False, isUnsafe := false }))\n\naddax\n\ntheorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := by\n  exact evil.elim\n")
+_add("macro_cheat_tactic", "env_manipulation", "A",
+     "macro \"cheat\" : tactic => `(tactic| exact sorryAx _ false)\n\ntheorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n  cheat\n")
+_add("macro_rules_override", "env_manipulation", "A",
+     "macro_rules | `(tactic| nlinarith $[[$_,*]]?) => `(tactic| exact sorryAx _ false)\n\ntheorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n  nlinarith [sq_nonneg (a - b)]\n")
+_add("run_cmd_adds_axiom", "env_manipulation", "C",
+     "open Lean Elab Command in\nrun_cmd liftCoreM (addDecl (.axiomDecl { name := `evil2, levelParams := [], type := mkConst ``False, isUnsafe := false }))\n\ntheorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := by\n  exact evil2.elim\n")
+_add("initialize_hook", "env_manipulation", "A",
+     "initialize myRef : IO.Ref Nat ← IO.mkRef 0\n\ntheorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n  have hb : b = 2 - a := by linarith\n  subst hb\n  nlinarith [sq_nonneg (a - 1)]\n")
+_add("syntax_hijack", "env_manipulation", "A",
+     "syntax \"my_tac\" : tactic\nmacro_rules | `(tactic| my_tac) => `(tactic| exact sorryAx _ false)\n\ntheorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n  my_tac\n")
+_add("skipKernelTC_option", "env_manipulation", "D",
+     "set_option debug.skipKernelTC true in\ntheorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n  norm_num\n",
+     note="the option does not exist on the pinned Lean 4.9.0-rc1 (`unknown option`), so this case is inert there; it guards a pin upgrade")
+_add("eval_side_effect", "env_manipulation", "D",
+     "#eval IO.println \"side effect\"\n\ntheorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n  norm_num\n")
+_add("exit_after_theorem", "env_manipulation", "D",
+     "theorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n  norm_num\n\n#exit\n\ntheorem unreachable : False := by sorry\n")
+# ---- theorem proved for a different name ---------------------------------------------
+_add("wrong_name_prime", "wrong_name", "A",
+     "theorem adv_a' (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK)
+_add("wrong_name_other", "wrong_name", "B",
+     "theorem adv_b_proof (n : ℕ) (h : n % 2 = 1) : n ^ 2 % 2 = 1 := by\n" + _B_OK)
+_add("namespace_wrapped", "wrong_name", "A",
+     "namespace Foo\ntheorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK + "end Foo\n")
+_add("private_target", "wrong_name", "A",
+     "private theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK)
+_add("attribute_on_target", "wrong_name", "A",
+     "@[simp] theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK)
+_add("duplicate_target", "wrong_name", "A",
+     "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK + "\ntheorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n  sorry\n")
+# ---- timeout / crash / memory bombs (own 1-worker pool, 30 s, 1 GiB) ------------------
+_add("timeout_bomb_norm_num_prime", "timeout_bomb", "P",
+     "theorem adv_p : (2:ℕ) + 2 = 4 := by\n  have h : Nat.Prime (2 ^ 127 - 1) := by norm_num\n  norm_num\n",
+     pool="bomb", note="Mersenne-127 primality by trial division: minutes, no heartbeat limit (maxHeartbeats 0)")
+_add("crash_stack_overflow", "timeout_bomb", "P",
+     "theorem adv_p : (2:ℕ) + 2 = 4 := by\n  have h : (List.range 20000).foldl (· + ·) 0 = 199990000 := by\n    set_option maxRecDepth 16000 in decide\n  norm_num\n",
+     pool="bomb", note="aborts the REPL process (lean::stack_space_exception)")
+_add("memory_bomb_pow", "timeout_bomb", "P",
+     "theorem adv_p : (2:ℕ) + 2 = 4 := by\n  have h : (3:ℕ) ^ (2 ^ 31 - 1) % 7 = (3:ℕ) ^ (2 ^ 31 - 1) % 7 := by decide\n  norm_num\n",
+     pool="bomb", note="multi-GB bignum; killed by the anon-memory watchdog")
+# ---- banned constructs (policy; the proof of the statement itself is valid) ------------
+_add("unsafe_implemented_by", "banned_construct", "D",
+     "unsafe def fastImpl (n : ℕ) : ℕ := n\n\n@[implemented_by fastImpl] def slowSpec (n : ℕ) : ℕ := n\n\ntheorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n  norm_num\n")
+_add("partial_def", "banned_construct", "D",
+     "partial def loop (n : ℕ) : ℕ := loop n\n\ntheorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n  norm_num\n")
+_add("extern_attr", "banned_construct", "D",
+     "@[extern \"c_fn\"] opaque cFn : ℕ → ℕ\n\ntheorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n  norm_num\n")
+_add("search_tactic_exact_q", "banned_construct", "A",
+     "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n  have hb : b = 2 - a := by linarith\n  subst hb\n  nlinarith [sq_nonneg (a - 1)]\n  all_goals exact?\n")
+_add("maxRecDepth_abuse", "banned_construct", "D",
+     "set_option maxRecDepth 100000000 in\ntheorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n  norm_num\n")
+_add("helper_def_with_proof", "banned_construct", "D",
+     "def twenty : ℕ := 20\n\ntheorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n  norm_num\n",
+     note="allow_helper_defs is off by default: a def is a policy reject")
+# ---- header tampering ----------------------------------------------------------------
+_add("header_dropped_open", "header_tamper", "A",
+     "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK,
+     raw_header="import Mathlib\nimport Aesop\n\nset_option maxHeartbeats 0\n\n")
+_add("header_extra_import", "header_tamper", "A",
+     "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK,
+     raw_header="import Mathlib\nimport Aesop\nimport Lean\n\nset_option maxHeartbeats 0\n\nopen BigOperators Real Nat Topology Rat\n\n")
+_add("header_heartbeats_changed", "header_tamper", "A",
+     "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK,
+     raw_header="import Mathlib\nimport Aesop\n\nset_option maxHeartbeats 400000\n\nopen BigOperators Real Nat Topology Rat\n\n")
+_add("header_missing", "header_tamper", "A",
+     "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK, raw_header="")
+# ---- no code -------------------------------------------------------------------------
+_add("empty_completion", "no_code", "A", None, completion="")
+_add("prose_only", "no_code", "A", None, completion="The statement follows from AM-GM; I leave the formalisation to the reader.")
+_add("unclosed_fence_truncated", "no_code", "A", None,
+     completion="plan...\n```lean4\n" + HEADER + "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK)
+_add("header_only", "no_code", "A", "", note="fence with only the header: no target theorem")
+
+# ======================================================================================
+# Known-good, hand-restated Mathlib-style theorems with proofs. Each is checked by the
+# independent compile-only pass BEFORE scoring; any that does not compile under the pin is
+# dropped and reported, not repaired after the fact.
+RESTATED: list[dict] = []
+
+
+def _rs(id, code):
+    RESTATED.append(dict(id=id, code=code))
+
+
+_rs("rs_add_comm", "theorem rs_add_comm (a b : ℕ) : a + b = b + a := by\n  omega\n")
+_rs("rs_sq_nonneg", "theorem rs_sq_nonneg (x : ℝ) : 0 ≤ x ^ 2 := by\n  positivity\n")
+_rs("rs_gauss", "theorem rs_gauss (n : ℕ) : (∑ i in Finset.range (n + 1), i) * 2 = n * (n + 1) := by\n  induction n with\n  | zero => simp\n  | succ k ih =>\n    rw [Finset.sum_range_succ, add_mul, ih]\n    ring\n")
+_rs("rs_amgm", "theorem rs_amgm (a b : ℝ) : a * b ≤ (a ^ 2 + b ^ 2) / 2 := by\n  nlinarith [sq_nonneg (a - b)]\n")
+_rs("rs_dvd", "theorem rs_dvd (n : ℕ) : 2 ∣ n * (n + 1) :=\n  even_iff_two_dvd.mp (Nat.even_mul_succ_self n)\n")
+_rs("rs_abs", "theorem rs_abs (x y : ℝ) : |x + y| ≤ |x| + |y| := abs_add x y\n")
+_rs("rs_odd_sum", "theorem rs_odd_sum (n : ℕ) : ∑ i in Finset.range n, (2 * i + 1) = n ^ 2 := by\n  induction n with\n  | zero => simp\n  | succ k ih => rw [Finset.sum_range_succ, ih]; ring\n")
+_rs("rs_exists", "theorem rs_exists : ∃ x : ℕ, x * x = 49 := ⟨7, by norm_num⟩\n")
+_rs("rs_irrational_nat", "theorem rs_irr : ¬ ∃ x : ℕ, x * x = 2 := by\n  rintro ⟨x, hx⟩\n  have h1 : x ≤ 1 := by nlinarith\n  interval_cases x <;> omega\n")
+_rs("rs_calc", "theorem rs_calc (a b : ℕ) (h : a ≤ b) : a + 1 ≤ b + 1 := by\n  calc a + 1 ≤ b + 1 := Nat.succ_le_succ h\n")
+_rs("rs_let_in_type", "theorem rs_let : let x := (5 : ℕ); x + x = 10 := by\n  intro x\n  rfl\n")
+_rs("rs_comm_ring", "theorem rs_comm_ring {α : Type*} [CommRing α] (a b : α) : (a + b) ^ 2 = a ^ 2 + 2 * a * b + b ^ 2 := by\n  ring\n")
+_rs("rs_lemma_kw", "lemma rs_lemma_kw (a b : ℚ) (h : a < b) : a < (a + b) / 2 := by\n  linarith\n")
+_rs("rs_set_eq", "theorem rs_set_eq : {x : ℕ | x < 3} = {0, 1, 2} := by\n  ext x\n  simp only [Set.mem_setOf_eq, Set.mem_insert_iff, Set.mem_singleton_iff]\n  omega\n")
+_rs("rs_card_filter", "theorem rs_card : (Finset.filter (fun n => n % 3 = 0) (Finset.range 30)).card = 10 := by\n  decide\n")
+_rs("rs_word_in_comment_and_string", "theorem rs_str (n : ℕ) : n + 0 = n := by\n  -- we never use sorry or native_decide here\n  have s : String := \"sorry /- not a comment\"\n  simp\n")
+_rs("rs_nested_comment", "theorem rs_nested (n : ℕ) : 0 + n = n := by\n  /- outer /- inner sorry -/ still outer admit -/\n  simp\n")
+_rs("rs_char_quote", "theorem rs_char (n : ℕ) : n * 1 = n := by\n  have c : Char := '\"'\n  simp\n")
+_rs("rs_set_option_in", "theorem rs_set_opt (a b c : ℝ) (h : a ≤ b) (hc : 0 ≤ c) : a * c ≤ b * c := by\n  set_option maxHeartbeats 400000 in\n  nlinarith [mul_nonneg (sub_nonneg.mpr h) hc]\n")
+_rs("rs_rcases", "theorem rs_rcases (n : ℕ) : n = 0 ∨ 0 < n := by\n  rcases n with _ | n\n  · left; rfl\n  · right; exact Nat.succ_pos n\n")
+_rs("rs_mod4", "theorem rs_mod4 (n : ℕ) : n % 4 = 0 ∨ n % 4 = 1 ∨ n % 4 = 2 ∨ n % 4 = 3 := by\n  omega\n")
+_rs("rs_prime97", "theorem rs_prime97 : Nat.Prime 97 := by\n  norm_num\n")
+_rs("rs_helper_lemma", "lemma aux_pos (x : ℝ) : 0 < x ^ 2 + 1 := by positivity\n\ntheorem rs_helper (x : ℝ) : 0 < (x ^ 2 + 1) * 3 := by\n  have := aux_pos x\n  linarith\n")
+_rs("rs_exp_pos", "theorem rs_exp_pos (x : ℝ) : 0 < Real.exp x := Real.exp_pos x\n")
+_rs("rs_by_contra", "theorem rs_by_contra (x : ℝ) (h : x ^ 2 < 0) : False := by\n  nlinarith [sq_nonneg x]\n")
+_rs("rs_em", "theorem rs_em (p : Prop) : p ∨ ¬ p := Classical.em p\n")
+_rs("rs_funext", "theorem rs_funext : (fun x : ℕ => x + 0) = (fun x => x) := by\n  funext x\n  simp\n")
+_rs("rs_field_simp", "theorem rs_div (a b : ℝ) (hb : 0 < b) : a / b * b = a := by\n  field_simp\n")
+_rs("rs_default_arg", "theorem rs_default_arg (n : ℕ := 5) : n + 0 = n := by\n  omega\n")
+_rs("rs_structure_instance", "theorem rs_prod (p : ℕ × ℕ) (h : p = { fst := 1, snd := 2 }) : p.1 + p.2 = 3 := by\n  subst h\n  rfl\n")
+
+
+def table_hash() -> str:
+    """SHA-256 of the canonical JSON of both case tables + the header + statements."""
+    blob = json.dumps(
+        {"header": HEADER, "statements": STATEMENTS, "adversarial": ADVERSARIAL, "restated": RESTATED},
+        sort_keys=True, ensure_ascii=False,
+    )
+    return hashlib.sha256(blob.encode("utf-8")).hexdigest()
+
+
+if __name__ == "__main__":
+    print(len(ADVERSARIAL), "adversarial,", len(RESTATED), "restated; sha256", table_hash())

code/scripts/lean_ceiling_test.py

@@ -0,0 +1,729 @@
+"""Ceiling test + measurements for the strict Lean verifier (Track L / P1).
+
+House rule (paid for twice: a perfect policy scored 0/100 through a broken extractor, and
+mprm's PPO runs measured the grader): ceiling-test every instrument before trusting any
+number it produces. For a VERIFIER the ceiling has two sides, and both are checked:
+
+  * known-good  : proofs that independently compile under the pinned toolchain must be
+                  ACCEPTED (>= 98%): a strict gate that rejects honest proofs starves RL.
+  * adversarial : hand-built exploits must ALL be REJECTED (0% accepted): a lax gate is
+                  exactly what a policy under RL finds.
+
+The criteria, the case tables (frozen, SHA-256 pinned) and the pass thresholds are
+pre-registered in docs/LEAN_VERIFIER.md BEFORE the registered run. `--verify-prereg`
+checks that the tables in this checkout hash to the registered value; the run refuses to
+start otherwise (override: --allow-unregistered, recorded in the output).
+
+Known-good sources
+  1. Goedel-LM/Lean-workbook-proofs (public, proofs verified by Goedel-Prover-V1 under
+     the same Lean 4.9 / Mathlib environment as DeepSeek-Prover-V1.5 and
+     Goedel-Prover-V2): 6 contiguous slices of 25 rows at fixed offsets, fetched over
+     HTTP from the HF datasets-server (never the full dataset).
+  2. 30 hand-restated Mathlib-style theorems with proofs (scripts/lean_ceiling_cases.py).
+  Each is first run through an INDEPENDENT compile-only pass (same REPL, no strictness at
+  all: no error message, no `sorry` warning). Items that fail it are version-incompatible
+  with the pin and are excluded and listed; the strict verdict is scored on the rest.
+
+Host rules (/apps/CLAUDE.md): `apply_resource_guards(heavy=True)` runs the preflight and
+takes the exclusive heavy-job flock; scratch lives under align/.tmp/; nothing touches / or
+/tmp; every blocking call has a timeout; the process has a global wall-clock watchdog.
+
+    python scripts/lean_ceiling_test.py --workers 8
+    python scripts/lean_ceiling_test.py --verify-prereg
+    python scripts/lean_ceiling_test.py --measure-only
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import math
+import os
+import re
+import statistics
+import subprocess
+import sys
+import threading
+import time
+import urllib.parse
+import urllib.request
+from pathlib import Path
+
+REPO = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(REPO))
+
+import importlib.util  # noqa: E402
+
+# Loaded by path (not `import lean_ceiling_cases`): scripts/ is not a package, and the
+# dependency firewall (rightly) rejects an import it cannot place in stdlib/allowed/align.
+_spec = importlib.util.spec_from_file_location("lean_ceiling_cases", Path(__file__).with_name("lean_ceiling_cases.py"))
+C = importlib.util.module_from_spec(_spec)
+sys.modules["lean_ceiling_cases"] = C
+_spec.loader.exec_module(C)
+
+DOC = REPO / "docs" / "LEAN_VERIFIER.md"
+DATASET = "Goedel-LM/Lean-workbook-proofs"
+SLICE_OFFSETS = (0, 5000, 10000, 15000, 20000, 25000)
+SLICE_LEN = 25
+KNOWN_GOOD_FLOOR = 0.98
+MIN_KNOWN_GOOD = 50
+MIN_ADVERSARIAL = 30
+
+
+# ---------------------------------------------------------------------------- prereg
+
+
+def registered_hash() -> str | None:
+    if not DOC.exists():
+        return None
+    m = re.search(r"PREREG_TABLE_SHA256:\s*`?([0-9a-f]{64})`?", DOC.read_text())
+    return m.group(1) if m else None
+
+
+def verify_prereg() -> bool:
+    reg, cur = registered_hash(), C.table_hash()
+    print(f"registered : {reg}\ncurrent    : {cur}\nmatch      : {reg == cur}")
+    return reg == cur
+
+
+# ---------------------------------------------------------------------------- verdict
+
+
+def empty_known_good_entry(candidates: int, excluded: int) -> dict:
+    """Report entry for a run in which NO known-good proof reached the strict stage (toolchain broken,
+    REPL crashing, header mismatch: every candidate failed the independent compile-only filter).
+    That is the absence of the ceiling's first side, i.e. a FAILURE: a verifier that rejects
+    everything would otherwise sail through the adversarial side (accepted == 0 is trivially true)."""
+    return {"n": 0, "accepted": 0, "rate": float("nan"), "floor": KNOWN_GOOD_FLOOR, "pass": False,
+            "why": f"no known-good proof survived the compile-only filter ({excluded}/{candidates} excluded)",
+            "rejected": []}
+
+
+def ceiling_ok(report: dict, only: str | None) -> bool:
+    """Exit-status verdict. Both sides must be PRESENT and passing in a full run; `--only` runs need
+    the side they ran. A missing side is never a pass (the old `all(v.get('pass', True) ...)` read
+    absence as success)."""
+    need = {None: ("known_good", "adversarial"), "good": ("known_good",), "adversarial": ("adversarial",)}[only]
+    if any(not isinstance(report.get(k), dict) or not report[k].get("pass") for k in need):
+        return False
+    return all(v.get("pass", True) for v in report.values() if isinstance(v, dict))
+
+
+# ---------------------------------------------------------------------------- statistics
+
+
+def wilson(k: int, n: int, z: float = 1.96) -> tuple[float, float]:
+    """Wilson score interval for a binomial proportion (exact-ish at k=n or k=0)."""
+    if n == 0:
+        return (float("nan"), float("nan"))
+    p = k / n
+    d = 1 + z * z / n
+    c = (p + z * z / (2 * n)) / d
+    h = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / d
+    return (max(0.0, c - h), min(1.0, c + h))
+
+
+def pctl(xs: list[float], q: float) -> float:
+    if not xs:
+        return float("nan")
+    s = sorted(xs)
+    return s[min(len(s) - 1, int(q * len(s)))]
+
+
+# ---------------------------------------------------------------------------- data
+
+
+def http_json(url: str, timeout: float = 60.0, retries: int = 4) -> dict:
+    last = None
+    for i in range(retries):
+        try:
+            req = urllib.request.Request(url, headers={"User-Agent": "align-lean-ceiling/1"})
+            with urllib.request.urlopen(req, timeout=timeout) as r:
+                return json.loads(r.read().decode("utf-8"))
+        except Exception as e:  # noqa: BLE001
+            last = e
+            time.sleep(2 * (i + 1))
+    raise RuntimeError(f"GET {url} failed: {last!r}")
+
+
+def fetch_known_good(out_dir: Path) -> tuple[list[dict], dict]:
+    rows: list[dict] = []
+    for off in SLICE_OFFSETS:
+        q = urllib.parse.urlencode(
+            {"dataset": DATASET, "config": "default", "split": "train", "offset": off, "length": SLICE_LEN}
+        )
+        d = http_json(f"https://datasets-server.huggingface.co/rows?{q}")
+        for r in d["rows"]:
+            rows.append({"problem_id": r["row"]["problem_id"], "full_proof": r["row"]["full_proof"], "row_idx": r["row_idx"]})
+    raw = out_dir / "known_good_raw.jsonl"
+    with open(raw, "w") as f:
+        for r in rows:
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")
+    meta = {"dataset": DATASET, "offsets": list(SLICE_OFFSETS), "slice_len": SLICE_LEN, "n_rows": len(rows),
+            "raw_sha256": hashlib.sha256(raw.read_bytes()).hexdigest()}
+    try:
+        info = http_json(f"https://huggingface.co/api/datasets/{DATASET}")
+        meta["dataset_revision"] = info.get("sha")
+    except Exception as e:  # noqa: BLE001
+        meta["dataset_revision"] = f"unavailable: {e!r}"
+    return rows, meta
+
+
+def derive_statement(code: str):
+    """(statement, target_name) from a full proof file: header + the LAST theorem/lemma's
+    signature (cut at the signature-ending `:=`, via the verifier's own scanner) + `sorry`."""
+    from align.verify.lean.strict import find_sig_end, lex
+
+    lx = lex(code)
+    toks = lx.toks
+    if lx.broken:
+        return None
+    first = next((i for i, t in enumerate(toks) if t.kind == "id" and t.text in
+                  ("theorem", "lemma", "def", "abbrev", "instance", "example", "structure", "inductive")), None)
+    thm = [i for i, t in enumerate(toks) if t.kind == "id" and t.text in ("theorem", "lemma")]
+    if first is None or not thm:
+        return None
+    k = thm[-1]
+    e = find_sig_end(toks, k + 2)
+    if e is None or k + 1 >= len(toks):
+        return None
+    name = toks[k + 1].text
+    header = code[: toks[first - 1].end] if first > 0 else ""
+    sig = code[toks[k + 2].start : toks[e].start]
+    return header + "\n\n" + f"theorem {name} {sig} := by\n  sorry\n", name
+
+
+def as_completion(code: str) -> str:
+    return "I will state a plan and then give the proof.\n\n```lean4\n" + code.rstrip("\n") + "\n```\n"
+
+
+# ---------------------------------------------------------------------------- compile-only
+
+
+def compile_only(w, statement: str, code: str, timeout: float = 120.0) -> dict:
+    """The NAIVE grader: elaborate the candidate and look at nothing but error messages.
+    Used (a) as the independent reference compile that defines "known-good", and (b) to show
+    how many adversarial cases a compile-only reward would pay."""
+    from align.verify.lean.repl_pool import ReplError
+    from align.verify.lean.strict import lex, parse_statement
+
+    ref = parse_statement(statement)
+    if not ref.ok:
+        return {"ok": False, "why": "unsupported_statement"}
+    lx = lex(code)
+    nh = len(ref.header_toks)
+    if lx.broken or tuple(t.text for t in lx.toks[:nh]) != ref.header_toks:
+        return {"ok": False, "why": "header"}
+    body = code[lx.toks[nh - 1].end :] if nh else code
+    try:
+        env = w.header_env(ref.header_text)
+        resp = w.command(body, env, timeout)
+    except ReplError as e:
+        return {"ok": False, "why": e.code}
+    msgs = [m for m in resp.get("messages", []) if isinstance(m, dict)]
+    err = [m for m in msgs if m.get("severity") == "error"]
+    sorry = bool(resp.get("sorries")) or any("sorry" in str(m.get("data", "")) for m in msgs)
+    return {"ok": not err, "clean": (not err) and (not sorry), "why": "error" if err else ("sorry" if sorry else "")}
+
+
+def naive_fresh(w, code: str, timeout: float = 120.0) -> dict:
+    """Compile-only on the candidate's OWN header (fresh import, no cache) -- for the
+    adversarial cases whose header is the attack. Each case gets a FRESH process: a second
+    `import Mathlib` in a live REPL costs +3.9GiB anonymous (see `split_imports`)."""
+    from align.verify.lean.repl_pool import ReplError
+
+    w.discard()
+    try:
+        proc = w.ensure()
+        resp = proc.request({"cmd": code}, timeout)
+    except ReplError as e:
+        w.discard()
+        return {"ok": False, "why": e.code}
+    w.discard()
+    err = [m for m in resp.get("messages", []) if isinstance(m, dict) and m.get("severity") == "error"]
+    return {"ok": not err, "why": "error" if err else ""}
+
+
+# ---------------------------------------------------------------------------- measurement
+
+
+def dir_bytes(p: Path) -> int:
+    out = subprocess.run(["du", "-sb", str(p)], capture_output=True, text=True, timeout=300)
+    return int(out.stdout.split()[0]) if out.returncode == 0 and out.stdout else -1
+
+
+def evict_page_cache(tc) -> dict:
+    """Drop the page cache for every .olean/.ilean under the install (POSIX_FADV_DONTNEED
+    works unprivileged for clean, unmapped pages) so the next import is a COLD read."""
+    n = tot = 0
+    roots = [tc.workspace / ".lake" / "packages", tc.root / "elan" / "toolchains"]
+    for r in roots:
+        for dp, _dn, fns in os.walk(r):
+            for fn in fns:
+                if fn.endswith((".olean", ".ilean", ".so")):
+                    fp = os.path.join(dp, fn)
+                    try:
+                        fd = os.open(fp, os.O_RDONLY)
+                        try:
+                            tot += os.fstat(fd).st_size
+                            os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
+                            n += 1
+                        finally:
+                            os.close(fd)
+                    except OSError:
+                        pass
+    return {"files": n, "bytes": tot}
+
+
+def cached_fraction(tc) -> float | None:
+    """Fraction of the Mathlib .olean bytes resident in page cache, via `fincore` if present."""
+    import shutil
+
+    f = shutil.which("fincore")
+    if not f:
+        return None
+    lib = tc.workspace / ".lake" / "packages" / "mathlib" / ".lake" / "build" / "lib" / "Mathlib"
+    files = [str(p) for p in lib.rglob("*.olean")][:4000]
+    try:
+        out = subprocess.run([f, "--bytes", "--raw", "--noheadings", *files], capture_output=True, text=True, timeout=120)
+        res = tot = 0
+        for line in out.stdout.splitlines():
+            parts = line.split()
+            if len(parts) >= 3 and parts[0].isdigit():
+                res += int(parts[0])
+                tot += int(parts[1])
+        return res / tot if tot else None
+    except Exception:  # noqa: BLE001
+        return None
+
+
+def measure_install(tc, workers: int, out: dict) -> None:
+    from align.verify.lean.repl_pool import ReplLimits, ReplPool, ReplProcess
+
+    m: dict = {}
+    m["disk_bytes"] = {
+        "total": dir_bytes(tc.root),
+        "elan_toolchain": dir_bytes(tc.root / "elan"),
+        "workspace": dir_bytes(tc.workspace),
+        "mathlib_package": dir_bytes(tc.workspace / ".lake" / "packages" / "mathlib"),
+        "repl": dir_bytes(tc.root / "repl"),
+    }
+    hdr = C.HEADER.rstrip("\n")
+
+    def one_import() -> tuple[float, dict | None]:
+        p = ReplProcess(tc.argv(), tc.child_env, cwd=tc.cwd, limits=ReplLimits())
+        p.start()
+        t = time.monotonic()
+        r = p.request({"cmd": hdr}, 900)
+        dt = time.monotonic() - t
+        assert "env" in r and not r.get("messages"), r
+        mem = p.memory()
+        p.kill()
+        return dt, mem
+
+    ev = evict_page_cache(tc)
+    m["evicted"] = ev
+    m["cached_fraction_after_evict"] = cached_fraction(tc)
+    cold, cold_mem = one_import()
+    m["cached_fraction_after_cold_import"] = cached_fraction(tc)
+    warm = [one_import()[0] for _ in range(3)]
+    m["import_s"] = {"cold": cold, "warm": warm, "warm_median": statistics.median(warm)}
+    m["single_repl_after_import_mib"] = cold_mem
+
+    # N workers concurrently warm: per-process PSS (RSS double-counts the shared oleans)
+    pool = ReplPool.from_toolchain(tc, n_workers=workers, scratch_dir=str(tc.root / "tmp" / "measure"))
+    try:
+        t = time.monotonic()
+        errs = pool.warm(hdr)
+        m["pool_warm_s"] = time.monotonic() - t
+        assert not any(errs), errs
+        rep = pool.memory_report()
+        m["n_workers"] = workers
+        m["per_worker_mib"] = rep
+        pss = [r["pss"] for r in rep if r]
+        m["pss_mib"] = {"min": min(pss), "mean": statistics.mean(pss), "max": max(pss), "sum": sum(pss)}
+        m["pss_anon_mib"] = {"mean": statistics.mean(r["pss_anon"] for r in rep if r)}
+        m["rss_sum_mib"] = sum(r["rss"] for r in rep if r)
+    finally:
+        pool.close()
+    out["install"] = m
+
+
+# ---------------------------------------------------------------------------- confusion
+
+
+def confusion(good_acc: list[bool], bad_acc: list[bool]) -> dict:
+    tp = sum(good_acc)
+    fn = len(good_acc) - tp
+    fp = sum(bad_acc)
+    tn = len(bad_acc) - fp
+    return {"TP": tp, "FN": fn, "FP": fp, "TN": tn, "n_good": len(good_acc), "n_bad": len(bad_acc)}
+
+
+def render_confusion(c: dict) -> str:
+    return (
+        "                      verdict ACCEPT   verdict REJECT   total\n"
+        f"truth GOOD proofs     {c['TP']:>14}   {c['FN']:>14}   {c['n_good']:>5}\n"
+        f"truth BAD  proofs     {c['FP']:>14}   {c['TN']:>14}   {c['n_bad']:>5}\n"
+    )
+
+
+# ---------------------------------------------------------------------------- main
+
+
+def sample_memory(pool, stop: threading.Event, rows: list) -> None:
+    while not stop.is_set():
+        try:
+            rep = pool.memory_report()
+            with open("/proc/meminfo") as f:
+                avail = next(int(l.split()[1]) for l in f if l.startswith("MemAvailable:")) / 1024**2
+            rows.append({"t": time.time(), "workers": rep, "mem_available_gb": avail})
+        except Exception:  # noqa: BLE001
+            pass
+        stop.wait(0.5)
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--workers", type=int, default=None)
+    ap.add_argument("--out", default=None)
+    ap.add_argument("--verify-prereg", action="store_true")
+    ap.add_argument("--allow-unregistered", action="store_true")
+    ap.add_argument("--measure-only", action="store_true")
+    ap.add_argument("--skip-measure", action="store_true")
+    ap.add_argument("--max-wall-s", type=float, default=3600.0)
+    ap.add_argument("--no-semantic-pass", action="store_true")
+    ap.add_argument("--no-naive-pass", action="store_true")
+    ap.add_argument("--only", choices=("good", "adversarial"), default=None)
+    ap.add_argument("--exploratory", action="store_true",
+                    help="debug run: <=2 workers, no heavy-job lock (a registered run always takes it)")
+    ap.add_argument("--lock-wait-s", type=float, default=7200.0)
+    ap.add_argument("--no-kernel-recheck", action="store_true",
+                    help="known-good pass without the kernel recheck (default: on, as in the shipped config)")
+    ap.add_argument("--kg-limit", type=int, default=0,
+                    help="exploratory only: keep the first N rows of each fetched slice (0 = all)")
+    args = ap.parse_args()
+
+    if args.verify_prereg:
+        return 0 if verify_prereg() else 2
+    registered = registered_hash() == C.table_hash()
+    if not registered and not args.allow_unregistered and not args.measure_only:
+        print("REFUSING: case tables do not match the registered hash (see --verify-prereg); "
+              "pass --allow-unregistered for an exploratory run", file=sys.stderr)
+        return 2
+
+    os.environ.setdefault("HF_HOME", "/apps/development/rlhf/hf_cache")
+    os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
+    from align.utils.guard import apply_resource_guards
+
+    apply_resource_guards(cpu_threads=8, heavy=False, label="lean_ceiling_test")
+    if not args.exploratory:
+        # A Lean pool is a heavy job (research report, section 7A): same preflight + the same
+        # exclusive flock as the trainers. If a GPU job holds it, WAIT -- never overlap.
+        from align.exit_codes import PreflightRefused
+        from align.utils.guard import acquire_heavy_job_lock, preflight_host_check
+
+        t_wait = time.monotonic()
+        while True:
+            try:
+                preflight_host_check(strict=True)
+                acquire_heavy_job_lock("lean_ceiling_test")
+                break
+            except PreflightRefused as e:
+                if time.monotonic() - t_wait > args.lock_wait_s:
+                    print(f"gave up waiting for the host: {e}", file=sys.stderr)
+                    return 83
+                print(f"host busy ({str(e)[:120]}...); retrying in 30s", flush=True)
+                time.sleep(30)
+    elif (args.workers or 2) > 2:
+        print("--exploratory caps workers at 2", file=sys.stderr)
+        args.workers = 2
+
+    from align.verify.lean import (
+        LeanAccuracy, ReplLimits, ReplPool, StrictConfig, StrictVerifier, default_workers, find_toolchain,
+    )
+    from align.verify.lean.repl_pool import meminfo_available_gb, plan_workers
+
+    tc = find_toolchain()
+    limits = ReplLimits()
+    n = args.workers or default_workers()
+    n = plan_workers(n, limits.anon_cap_gb, 4.0, meminfo_available_gb())
+    stamp = time.strftime("%Y%m%d_%H%M%S")
+    out_dir = Path(args.out or REPO / ".tmp" / "lean_ceiling" / stamp)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    pools: list = []
+
+    def watchdog() -> None:
+        time.sleep(args.max_wall_s)
+        print(f"WATCHDOG: {args.max_wall_s}s exceeded; killing pools", file=sys.stderr, flush=True)
+        for p in pools:
+            try:
+                p.close()
+            except Exception:  # noqa: BLE001
+                pass
+        os._exit(124)
+
+    threading.Thread(target=watchdog, daemon=True).start()
+
+    def on_signal(signum, _frame):
+        print(f"signal {signum}: closing pools", file=sys.stderr, flush=True)
+        for p in pools:
+            try:
+                p.close()
+            except Exception:  # noqa: BLE001
+                pass
+        os._exit(128 + signum)
+
+    import signal
+
+    signal.signal(signal.SIGTERM, on_signal)
+    signal.signal(signal.SIGINT, on_signal)
+
+    result: dict = {
+        "started": stamp, "toolchain": tc.describe(), "pins": tc.pins, "workers": n,
+        "limits": limits.__dict__, "prereg_registered": registered, "table_sha256": C.table_hash(),
+        "cores": os.cpu_count(),
+    }
+
+    def new_pool(workers: int, lim: ReplLimits = limits, label: str = "pool") -> ReplPool:
+        p = ReplPool.from_toolchain(tc, n_workers=workers, limits=lim, scratch_dir=str(tc.root / "tmp" / label))
+        pools.append(p)
+        return p
+
+    # -------------------------------------------------- install measurements
+    if not args.skip_measure and args.only is None:
+        print("== measuring install ==", flush=True)
+        measure_install(tc, n, result)
+        print(json.dumps(result["install"], indent=1, default=str)[:3000], flush=True)
+        if args.measure_only:
+            (out_dir / "result.json").write_text(json.dumps(result, indent=1, default=str))
+            return 0
+
+    # -------------------------------------------------- known-good
+    good_items: list[dict] = []
+    if args.only in (None, "good"):
+        print("== known-good: fetching ==", flush=True)
+        rows, meta = fetch_known_good(out_dir)
+        if args.kg_limit:
+            if not args.exploratory:
+                print("--kg-limit is for exploratory runs only", file=sys.stderr)
+                return 2
+            rows = [r for i, r in enumerate(rows) if i % SLICE_LEN < args.kg_limit]
+            meta["kg_limit"] = args.kg_limit
+        result["known_good_source"] = meta
+        cands = []
+        for r in rows:
+            code = r["full_proof"]
+            ds = derive_statement(code)
+            if ds is None:
+                cands.append({"id": r["problem_id"], "source": "lean-workbook-proofs", "code": code, "statement": None})
+            else:
+                cands.append({"id": r["problem_id"], "source": "lean-workbook-proofs", "code": code, "statement": ds[0]})
+        for rs in C.RESTATED:
+            code = C.HEADER + rs["code"]
+            ds = derive_statement(code)
+            cands.append({"id": rs["id"], "source": "restated", "code": code, "statement": ds[0] if ds else None})
+        result["known_good_candidates"] = len(cands)
+
+        print("== known-good: independent compile-only filter ==", flush=True)
+        pool = new_pool(n, label="kg_compile")
+        t0 = time.monotonic()
+        refs = pool.map(
+            lambda w, it: compile_only(w, it["statement"], it["code"]) if it["statement"] else {"ok": False, "clean": False, "why": "no_statement"},
+            cands,
+        )
+        pool.close()
+        excluded = []
+        for it, r in zip(cands, refs):
+            if isinstance(r, Exception):
+                r = {"ok": False, "clean": False, "why": repr(r)}
+            it["ref"] = r
+            if r.get("clean"):
+                good_items.append(it)
+            else:
+                excluded.append({"id": it["id"], "source": it["source"], "why": r.get("why")})
+        result["known_good_excluded"] = excluded
+        result["known_good_filter_s"] = time.monotonic() - t0
+        print(f"compile-clean: {len(good_items)}/{len(cands)}  excluded: {len(excluded)}", flush=True)
+
+    # -------------------------------------------------- known-good: strict
+    good_results: list[dict] = []
+    if good_items:
+        print("== known-good: strict verification ==", flush=True)
+        pool = new_pool(n, label="kg_strict")
+        acc = LeanAccuracy(StrictVerifier(pool, StrictConfig(kernel_recheck=not args.no_kernel_recheck)))
+        from align.verify.lean.strict import AUDIT_PRELUDE, static_check
+
+        plan, _, _ = static_check(good_items[0]["statement"], good_items[0]["code"])
+        t = time.monotonic()
+        pool.warm(plan.header if plan else C.HEADER, AUDIT_PRELUDE)
+        result["strict_warm_s"] = time.monotonic() - t
+        result["warm_memory_mib"] = pool.memory_report()
+        stop, mem_rows = threading.Event(), []
+        th = threading.Thread(target=sample_memory, args=(pool, stop, mem_rows), daemon=True)
+        th.start()
+        t = time.monotonic()
+        def prog(ix, r):
+            print(f"  [{ix:3d}] {good_items[ix]['id']:34s} {'ACCEPT' if r.accepted else 'REJECT'} {r.elapsed_s:7.2f}s {list(r.reasons)}", flush=True)
+
+        accs, res = acc.accuracy_batch([g["statement"] for g in good_items], [as_completion(g["code"]) for g in good_items], progress=prog)
+        wall = time.monotonic() - t
+        stop.set()
+        th.join(5)
+        for g, a, r in zip(good_items, accs, res):
+            good_results.append({"id": g["id"], "source": g["source"], "accepted": bool(a), "reasons": list(r.reasons),
+                                 "diag": r.diagnostics, "elapsed_s": r.elapsed_s})
+        lat = [r.elapsed_s for r in res]
+        result["throughput"] = {
+            "n_proofs": len(res), "wall_s": wall, "proofs_per_s": len(res) / wall, "workers": n,
+            "latency_s": {"p50": pctl(lat, 0.5), "p90": pctl(lat, 0.9), "p99": pctl(lat, 0.99), "max": max(lat), "mean": statistics.mean(lat)},
+            "note": "unique statement per proof (no reference-env reuse across proofs); warm workers",
+        }
+        # per-worker PSS during the run
+        per_w: dict[int, list] = {}
+        for row in mem_rows:
+            for i, m in enumerate(row["workers"]):
+                if m:
+                    per_w.setdefault(i, []).append(m)
+        result["pss_during_run_mib"] = {
+            str(i): {"pss_mean": statistics.mean(m["pss"] for m in ms), "pss_max": max(m["pss"] for m in ms),
+                     "pss_anon_max": max(m["pss_anon"] for m in ms), "rss_max": max(m["rss"] for m in ms),
+                     "samples": len(ms)} for i, ms in sorted(per_w.items())
+        }
+        result["host_during_run"] = {"min_mem_available_gb": min((r["mem_available_gb"] for r in mem_rows), default=None)}
+        result["pool_stats"] = pool.stats.__dict__
+        # repeat-G run: GRPO shape (same statement x8, reference env reused) on the 12
+        # MEDIAN-latency statements of the pass above (the heavy tail would swamp a 96-proof run)
+        order = sorted(range(len(good_items)), key=lambda i: res[i].elapsed_s)
+        mid = max(0, len(order) // 2 - 6)
+        rep = [good_items[i] for i in order[mid : mid + 12]] * 8
+        t = time.monotonic()
+        accs2, _ = acc.accuracy_batch([g["statement"] for g in rep], [as_completion(g["code"]) for g in rep])
+        w2 = time.monotonic() - t
+        result["throughput_grpo_shape"] = {"n_proofs": len(rep), "wall_s": w2, "proofs_per_s": len(rep) / w2, "accepted": int(sum(accs2)),
+                                           "shape": "12 median-latency statements x 8 identical samples"}
+        pool.close()
+
+    # -------------------------------------------------- adversarial
+    adv_results: list[dict] = []
+    if args.only in (None, "adversarial"):
+        print("== adversarial: strict ==", flush=True)
+
+        def completion_for(case) -> tuple[str, str | None]:
+            stmt = C.HEADER + C.STATEMENTS[case["statement"]]
+            if case["completion"] is not None:
+                return stmt, case["completion"]
+            hdr = C.HEADER if case["raw_header"] is None else case["raw_header"]
+            return stmt, as_completion(hdr + case["code"])
+
+        main_cases = [c for c in C.ADVERSARIAL if c["pool"] == "main"]
+        bomb_cases = [c for c in C.ADVERSARIAL if c["pool"] == "bomb"]
+        pool = new_pool(n, label="adv_strict")
+        acc = LeanAccuracy(StrictVerifier(pool, StrictConfig(proof_timeout_s=120)), min_batch=10**9)
+        pairs = [completion_for(c) for c in main_cases]
+        _, rs = acc.accuracy_batch([p[0] for p in pairs], [p[1] for p in pairs])
+        pool.close()
+        by_id = {}
+        for c, r in zip(main_cases, rs):
+            by_id[c["id"]] = {"id": c["id"], "category": c["category"], "accepted": r.accepted, "reasons": list(r.reasons),
+                              "stage": r.stage, "diag": r.diagnostics, "elapsed_s": r.elapsed_s, "pool": "main"}
+        # bombs: 1 worker, 1 GiB anon cap, 30 s budget so they terminate promptly (same code paths)
+        blim = ReplLimits(anon_cap_gb=1.0)
+        bpool = new_pool(1, blim, label="adv_bomb")
+        bacc = LeanAccuracy(StrictVerifier(bpool, StrictConfig(proof_timeout_s=30)), min_batch=10**9)
+        for c in bomb_cases:
+            stmt, comp = completion_for(c)
+            sc = bacc.score(stmt, comp)
+            by_id[c["id"]] = {"id": c["id"], "category": c["category"], "accepted": sc.result.accepted,
+                              "reasons": list(sc.result.reasons), "stage": sc.result.stage, "diag": sc.result.diagnostics,
+                              "elapsed_s": sc.result.elapsed_s, "pool": "bomb(1w,1GiB,30s)"}
+        bpool.close()
+        adv_results = [by_id[c["id"]] for c in C.ADVERSARIAL]
+
+        # naive compile-only on each case's own text (what a compile-only reward would pay)
+        if not args.no_naive_pass:
+            print("== adversarial: naive compile-only (liveness) ==", flush=True)
+            npool = new_pool(n, label="adv_naive")
+            live_cases = [c for c in main_cases if c["completion"] is None and c["code"]]
+            texts = [(c["raw_header"] if c["raw_header"] is not None else C.HEADER) + c["code"] for c in live_cases]
+            nv = npool.map(lambda w, t: naive_fresh(w, t, 120.0), texts)
+            npool.close()
+            naive = {c["id"]: (v if not isinstance(v, Exception) else {"ok": False, "why": repr(v)}) for c, v in zip(live_cases, nv)}
+            for r in adv_results:
+                r["naive_compile_ok"] = naive.get(r["id"], {}).get("ok")
+
+        # semantic layer alone (lexical gate OFF): how much does the Lean-side audit catch by itself?
+        if not args.no_semantic_pass:
+            print("== adversarial: semantic layer only (lexical gate off) ==", flush=True)
+            spool = new_pool(n, label="adv_sem")
+            sv = StrictVerifier(spool, StrictConfig(lexical_gate=False, kernel_recheck=True, proof_timeout_s=120))
+            sem_cases = [c for c in main_cases if c["completion"] is None and c["code"]]
+            from align.verify.lean.reward import extract_lean_code
+
+            sem_items = []
+            for c in sem_cases:
+                stmt, comp = completion_for(c)
+                sem_items.append((stmt, extract_lean_code(comp)))
+            sres = sv.verify_batch(sem_items)
+            spool.close()
+            sem = {c["id"]: (r.accepted, list(r.reasons)) for c, r in zip(sem_cases, sres)}
+            for r in adv_results:
+                if r["id"] in sem:
+                    r["semantic_only_accepted"], r["semantic_only_reasons"] = sem[r["id"]]
+
+    # -------------------------------------------------- report
+    report = {}
+    if not good_results and args.only in (None, "good"):
+        report["known_good"] = empty_known_good_entry(
+            result.get("known_good_candidates", 0), len(result.get("known_good_excluded", []))
+        )
+    if good_results:
+        acc_good = [g["accepted"] for g in good_results]
+        lo, hi = wilson(sum(acc_good), len(acc_good))
+        report["known_good"] = {"n": len(acc_good), "accepted": sum(acc_good), "rate": sum(acc_good) / len(acc_good),
+                                "wilson95": [lo, hi], "floor": KNOWN_GOOD_FLOOR,
+                                "pass": (sum(acc_good) / len(acc_good) >= KNOWN_GOOD_FLOOR and len(acc_good) >= MIN_KNOWN_GOOD),
+                                "rejected": [g for g in good_results if not g["accepted"]]}
+    if adv_results:
+        acc_bad = [a["accepted"] for a in adv_results]
+        by_cat: dict[str, list] = {}
+        for a in adv_results:
+            by_cat.setdefault(a["category"], []).append(a["accepted"])
+        report["adversarial"] = {"n": len(acc_bad), "accepted": sum(acc_bad), "pass": (sum(acc_bad) == 0 and len(acc_bad) >= MIN_ADVERSARIAL),
+                                 "by_category": {k: {"n": len(v), "accepted": sum(v)} for k, v in sorted(by_cat.items())},
+                                 "accepted_cases": [a for a in adv_results if a["accepted"]]}
+        live = [a for a in adv_results if a.get("naive_compile_ok")]
+        report["adversarial"]["naive_compile_only_would_accept"] = len(live)
+        if any("semantic_only_accepted" in a for a in adv_results):
+            sem_acc = [a for a in adv_results if a.get("semantic_only_accepted")]
+            report["adversarial"]["semantic_only_would_accept"] = [a["id"] for a in sem_acc]
+    if good_results and adv_results:
+        report["confusion"] = confusion([g["accepted"] for g in good_results], [a["accepted"] for a in adv_results])
+    result["report"] = report
+    result["good_results"] = good_results
+    result["adversarial_results"] = adv_results
+    (out_dir / "result.json").write_text(json.dumps(result, indent=1, default=str, ensure_ascii=False))
+    print("\n================ RESULT ================")
+    if "confusion" in report:
+        print(render_confusion(report["confusion"]))
+    if report.get("known_good", {}).get("n"):
+        k = report["known_good"]
+        print(f"known-good accepted {k['accepted']}/{k['n']} = {k['rate']:.4f}  Wilson95 [{k['wilson95'][0]:.4f}, {k['wilson95'][1]:.4f}]  floor {k['floor']}  PASS={k['pass']}")
+        for g in k["rejected"]:
+            print("   REJECTED known-good:", g["id"], g["reasons"], g["diag"])
+    if "adversarial" in report:
+        a = report["adversarial"]
+        print(f"adversarial accepted {a['accepted']}/{a['n']}  PASS={a['pass']}  (naive compile-only would accept {a['naive_compile_only_would_accept']})")
+        for c in a["accepted_cases"]:
+            print("   ACCEPTED adversarial:", c["id"], c["reasons"])
+    if "throughput" in result:
+        print("throughput:", json.dumps(result["throughput"]))
+    print(f"results: {out_dir/'result.json'}")
+    if "known_good" in report and not report["known_good"]["n"]:
+        print(f"known-good: NO EVIDENCE -- {report['known_good']['why']}  PASS=False")
+    return 0 if ceiling_ok(report, args.only) else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())