code/ar_differentiation_bed.py

"""ar_differentiation_bed.py — THE FOCUS (2026-07-09 redirect, the operator verbatim):
"refining the autoregressive techniques for differentiation rather than attempting
to just mash numbers together."

Differentiation is cultivated by PREDICTIVE pressure along the sequence — the
address parameterizing the next-byte distribution (Law 2: chain-rule advantage pays
ONLY where the composed address directly parameterizes the predictive distribution).
This bed puts the aleph in the autoregressive gradient path and measures what
differentiates. It is the Law-2 construction (codebook-pressure C3) + Tree 3d in
one harness; the Jun-19 "discuss before building" gate was resolved by the redirect.

Byte-level causal LM on wikitext-2-raw (HF parquet, CDN-fast), block 256. ARMS:
  sdpa       — standard causal transformer control (matched trunk).
  hub        — attention replaced by CAUSAL HUB: linear attention whose feature map
               is the 2K-oriented aleph address, prefix-sum memories (no selection
               event; O(n*K*d)). Differentiation cultivated INSIDE attention.
  addr_head  — sdpa trunk, but the OUTPUT HEAD reads ONLY the signed aleph
               coefficient vector w_k = sinh(u_k)/sum_j cosh(u_j) of the final
               hidden state (K -> 256 logits). The address MUST carry every bit of
               next-byte information — the hardest Law-2 bottleneck.

JUDGED BY: val bits-per-byte per arm (task) + CULTIVATION VITALS on every aleph
codebook (readouts, never losses): axis aliveness/hppl, drift-from-init +
binding fraction @0.29154, winner-|cos| saturation (sign-code emergence), shadow
path diversity (fixed high-bits hash). Never by recon.

Riders: pure Adam wd=0; no BN/Dropout/GAP on geometric paths; orthogonal init;
Colab-cell-safe (paste-ahead imports, no bare argparse, no __file__ reliance);
GPU-only for verdict runs; data_root OUTSIDE the mind repo.

Terminal:  python ar_differentiation_bed.py            # shapes/parse smoke
           python ar_differentiation_bed.py --train    # verdict run
Colab:     paste geolip_vitals.py cell, then this file (smoke auto-runs),
           then  train(steps=2000, data_root="/content/data")  in the next cell.
"""
from __future__ import annotations
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

if "anchor_drift" not in globals():
    try:
        from geolip_vitals import anchor_drift, axis_aliveness, path_diversity
    except ImportError:
        _here = globals().get("__file__")
        if _here is not None:
            import sys, pathlib
            sys.path.insert(0, str(pathlib.Path(_here).parent))
            from geolip_vitals import anchor_drift, axis_aliveness, path_diversity
        else:
            raise ImportError(
                "geolip_vitals not found — paste/run its cell first, or "
                "keep geolip_vitals.py beside this file (ships in this repo).")

VOCAB = 256  # bytes


# ------------------------------------------------------------------ aleph address
def _super_fibonacci_s3(n: int) -> torch.Tensor:
    """Near-uniform unit quaternions (Alexa CVPR'22; constants per canon) —
    starts the codebook INSIDE the RP^3 attractor basin. D=4 only."""
    PHI, PSI = math.sqrt(2.0), 1.533751168755204288118041
    i = torch.arange(n, dtype=torch.float64)
    s = (i + 0.5) / n
    r, R = torch.sqrt(s), torch.sqrt(1.0 - s)
    a, b = 2 * math.pi * i / PHI, 2 * math.pi * i / PSI
    q = torch.stack([r * torch.sin(a), r * torch.cos(a),
                     R * torch.sin(b), R * torch.cos(b)], dim=-1)
    return F.normalize(q, dim=-1).float()


class AlephAddress(nn.Module):
    """Closed-form aleph over 2K oriented half-axes (canon/aleph_core.md).
    signed(x): (..., K) w_k = sinh(u_k)/sum_j cosh(u_j)   — the Law-2 head feature.
    oriented(x): ((..., K), (..., K)) positive halves of the 2K softmax — HUB map."""

    def __init__(self, K: int, D: int, tau: float = 0.1, init: str = "random"):
        super().__init__()
        self.K, self.D, self.tau = K, D, tau
        if init == "fibonacci":
            assert D == 4, "fibonacci init lives on S^3 (D=4)"
            A = _super_fibonacci_s3(K)
        else:
            A = F.normalize(torch.randn(K, D), dim=-1)
        self.codebook = nn.Parameter(A)
        self.register_buffer("home", self.codebook.detach().clone())

    def _u(self, x):
        A = F.normalize(self.codebook, dim=-1)
        return (F.normalize(x, dim=-1) @ A.transpose(-1, -2)) / self.tau

    def oriented(self, x):
        u = self._u(x)
        m = u.abs().amax(dim=-1, keepdim=True)
        ep, en = torch.exp(u - m), torch.exp(-u - m)
        Z = (ep + en).sum(dim=-1, keepdim=True)
        return ep / Z, en / Z

    def signed(self, x):
        u = self._u(x)
        m = u.abs().amax(dim=-1, keepdim=True)
        ep, en = torch.exp(u - m), torch.exp(-u - m)
        return (ep - en) / (ep + en).sum(dim=-1, keepdim=True)

    def signed_at(self, x, taus):
        """Multi-tau stroboscope (rule of 3): signed coefficients at several
        temperatures, concatenated — softer taus keep the vector dense while a
        hard tau supplies the sign-code sharpness. v2 refinement (b)."""
        A = F.normalize(self.codebook, dim=-1)
        cos = F.normalize(x, dim=-1) @ A.transpose(-1, -2)
        outs = []
        for t in taus:
            u = cos / t
            m = u.abs().amax(dim=-1, keepdim=True)
            ep, en = torch.exp(u - m), torch.exp(-u - m)
            outs.append((ep - en) / (ep + en).sum(dim=-1, keepdim=True))
        return torch.cat(outs, dim=-1)

    def m_hat(self, x):
        """Closed-form soft read (decoders read M_hat, never M). v2 control (c)."""
        u = self._u(x)
        m = u.abs().amax(dim=-1, keepdim=True)
        ep, en = torch.exp(u - m), torch.exp(-u - m)
        A = F.normalize(self.codebook, dim=-1)
        return ((ep - en) @ A) / (ep + en).sum(dim=-1, keepdim=True)

    def m_hard_ste(self, x):
        """Canon hard mode: M_hard = sign(cos_win) * A[win], straight-through to
        the soft read — forward fully discrete SIGN CODE, backward soft gradient.
        Legal per theme A (reconstructive sign code, not a one-hot roster pick)."""
        u = self._u(x)
        soft = self.m_hat(x)
        win = u.abs().argmax(dim=-1)
        A = F.normalize(self.codebook, dim=-1)
        sign = torch.sign(torch.gather(u, -1, win.unsqueeze(-1))).squeeze(-1)
        hard = sign.unsqueeze(-1) * A[win]
        return hard + soft - soft.detach()

    @torch.no_grad()
    def vitals(self, x_sample) -> dict:
        u = self._u(x_sample.reshape(-1, x_sample.shape[-1]))
        p, n = self.oriented(x_sample.reshape(-1, x_sample.shape[-1]))
        two_k = torch.cat([p, n], dim=-1)
        win = two_k.argmax(dim=-1)
        cos_win = (u.abs().amax(dim=-1) * self.tau)  # winner |cos| — sign-code sat.
        d = anchor_drift(self.codebook, self.home)
        return {"drift": round(d["mean"], 4),
                "binding_frac": round(d["binding_fraction"], 4),
                "aliveness": axis_aliveness(two_k),
                "win_cos_mean": round(cos_win.mean().item(), 4),
                "paths": path_diversity(win)}


# ------------------------------------------------------------------------- blocks
class CausalSDPA(nn.Module):
    def __init__(self, d: int, heads: int = 4):
        super().__init__()
        self.h = heads
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        self.o = nn.Linear(d, d, bias=False)
        nn.init.orthogonal_(self.qkv.weight); nn.init.orthogonal_(self.o.weight)

    def forward(self, x):
        B, n, d = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q, k, v = (t.view(B, n, self.h, d // self.h).transpose(1, 2) for t in (q, k, v))
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.o(y.transpose(1, 2).reshape(B, n, d))


class CausalHUB(nn.Module):
    """Causal aleph linear attention: prefix-sum memories over the two K-wide
    halves of the oriented address; 2K never materialized; no selection event."""

    def __init__(self, d: int, K: int = 32, D: int = 4, tau: float = 0.1):
        super().__init__()
        self.addr = AlephAddress(K, D, tau)
        self.q = nn.Linear(d, D, bias=False)
        self.k = nn.Linear(d, D, bias=False)
        self.v = nn.Linear(d, d, bias=False)
        self.o = nn.Linear(d, d, bias=False)
        for m in (self.q, self.k, self.v, self.o):
            nn.init.orthogonal_(m.weight)

    def forward(self, x):
        qp, qn = self.addr.oriented(self.q(x))          # (B, n, K)
        kp, kn = self.addr.oriented(self.k(x))
        v = self.v(x)                                    # (B, n, d)
        Sp = torch.cumsum(torch.einsum("bnk,bnd->bnkd", kp, v), dim=1)
        Sn = torch.cumsum(torch.einsum("bnk,bnd->bnkd", kn, v), dim=1)
        zp = torch.cumsum(kp, dim=1)
        zn = torch.cumsum(kn, dim=1)
        num = torch.einsum("bnk,bnkd->bnd", qp, Sp) + torch.einsum("bnk,bnkd->bnd", qn, Sn)
        den = (qp * zp).sum(-1, keepdim=True) + (qn * zn).sum(-1, keepdim=True)
        return self.o(num / den.clamp_min(1e-12))


class MslRelay(nn.Module):
    """Depth-composition unit (chain-rule probe): multi-slot M_hat read entering
    the trunk as a NEAR-ZERO gated residual (gate init -3.0, sigma~0.047 — theme D:
    geometry enters as a nudge and grows only if it earns gradient)."""

    def __init__(self, d: int, n_slots: int = 16, K: int = 64):
        super().__init__()
        self.n_slots = n_slots
        self.proj = nn.Linear(d, n_slots * 4, bias=False)
        self.out = nn.Linear(n_slots * 4, d, bias=False)
        nn.init.orthogonal_(self.proj.weight)
        nn.init.orthogonal_(self.out.weight)
        self.addr = AlephAddress(K, 4)
        self.gate = nn.Parameter(torch.tensor(-3.0))

    def forward(self, x):
        B, n, _ = x.shape
        slots = self.proj(x).view(B, n, self.n_slots, 4)
        m = self.addr.m_hat(slots).reshape(B, n, -1)
        return x + self.gate.sigmoid() * self.out(m)


class Block(nn.Module):
    def __init__(self, d: int, attn: nn.Module):
        super().__init__()
        self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.attn = attn
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):
        x = x + self.attn(self.n1(x))
        return x + self.mlp(self.n2(x))


class ByteLM(nn.Module):
    def __init__(self, arm: str, d: int = 192, layers: int = 4, block: int = 256,
                 K: int = 32, D: int = 4):
        super().__init__()
        # "<arm>_tri" suffix = trigram byte embedding (AlephLM byte_emb x3 lineage):
        # token embedding is the sum of embeddings of bytes t, t-1, t-2.
        self.trigram = arm.endswith("_tri")
        if self.trigram:
            arm = arm[:-4]
        # "_fib" = super-Fibonacci S^3 codebook init (basin test: starts INSIDE
        # the RP^3 attractor; primary observable is init->final geodesic drift).
        self.fib = arm.endswith("_fib")
        if self.fib:
            arm = arm[:-4]
        # "relay*" = stacked addresses in depth: MslRelay after every block.
        # relay -> sdpa trunk + standard head; relay_msl64 -> + addressed head.
        self.use_relay = arm.startswith("relay")
        if arm == "relay":
            arm = "sdpa"
        elif arm == "relay_msl64":
            arm = "addr_msl64"
        self.arm, self.block = arm, block
        self.emb = nn.Embedding(VOCAB, d)
        if self.trigram:
            self.emb1 = nn.Embedding(VOCAB, d)
            self.emb2 = nn.Embedding(VOCAB, d)
        self.pos = nn.Parameter(torch.zeros(1, block, d) + 0.01 * torch.randn(1, block, d))
        mk_attn = (lambda: CausalHUB(d, K, D)) if arm == "hub" else (lambda: CausalSDPA(d))
        self.blocks = nn.ModuleList([Block(d, mk_attn()) for _ in range(layers)])
        if self.use_relay:
            self.relays = nn.ModuleList([MslRelay(d) for _ in range(layers)])
        self.nf = nn.LayerNorm(d)
        if arm == "addr_head":
            self.head_addr = AlephAddress(K, d)  # v1: codebook in model dim — COLLAPSED
            self.head = nn.Linear(K, VOCAB, bias=True)
        elif arm in ("addr_d4", "addr_3tau", "addr_mhat"):
            # v2 refinements: LOW-D HOME — learned projection to the canon D=4 home
            # before addressing (mirrors the healthy HUB arms), K=64.
            self.head_proj = nn.Linear(d, 4, bias=False)
            nn.init.orthogonal_(self.head_proj.weight)
            self.head_addr = AlephAddress(64, 4)
            if arm == "addr_d4":
                self.head = nn.Linear(64, VOCAB, bias=True)      # w alone, D=4 home
            elif arm == "addr_3tau":
                self.taus = (0.05, 0.1, 0.3)                     # rule-of-3 strobe
                self.head = nn.Linear(64 * 3, VOCAB, bias=True)
            else:                                                 # addr_mhat
                self.head = nn.Linear(4, VOCAB, bias=True)       # tightest: M_hat
        elif arm.startswith("addr_msl"):
            # v3: MULTI-SLOT heads — the 16s funnel widening: P parallel D=4 slots
            # over a SHARED codebook. addr_msl consumes the reconstructive M_hat per
            # slot (Px4 dims); addr_msl_w consumes signed w per slot (Px64) — tests
            # whether slot-parallel consumption alone rescues the coefficient path.
            # addr_msl<P> = slot-count dose-response. addr_mslh<P> = HARD sign-code
            # consumption (straight-through M_hard per slot).
            self.hard = arm.startswith("addr_mslh")
            if arm in ("addr_msl", "addr_msl_w"):
                self.n_slots = 16
            else:
                self.n_slots = int(arm[len("addr_mslh" if self.hard else "addr_msl"):])
            self.head_proj = nn.Linear(d, self.n_slots * 4, bias=False)
            nn.init.orthogonal_(self.head_proj.weight)
            self.head_addr = AlephAddress(
                64, 4, init="fibonacci" if self.fib else "random")
            width = self.n_slots * (64 if arm == "addr_msl_w" else 4)
            self.head = nn.Linear(width, VOCAB, bias=True)
        elif arm == "addr_3tau_mhat":
            # v3: combine the two v2 winners — 3-tau stroboscope + reconstructive read.
            self.head_proj = nn.Linear(d, 4, bias=False)
            nn.init.orthogonal_(self.head_proj.weight)
            self.head_addr = AlephAddress(64, 4)
            self.taus = (0.05, 0.1, 0.3)
            self.head = nn.Linear(64 * 3 + 4, VOCAB, bias=True)
        else:
            self.head = nn.Linear(d, VOCAB, bias=True)
        self._last_h = None

    def forward(self, idx):
        x = self.emb(idx)
        if self.trigram:  # past-only shifts — causality preserved
            x = x + self.emb1(F.pad(idx, (1, 0), value=0)[:, :-1]) \
                  + self.emb2(F.pad(idx, (2, 0), value=0)[:, :-2])
        x = x + self.pos[:, : idx.shape[1]]
        if self.use_relay:
            for b, r in zip(self.blocks, self.relays):
                x = r(b(x))
        else:
            for b in self.blocks:
                x = b(x)
        h = self.nf(x)
        self._last_h = h.detach()
        if self.arm == "addr_head":
            return self.head(self.head_addr.signed(h))
        if self.arm == "addr_d4":
            return self.head(self.head_addr.signed(self.head_proj(h)))
        if self.arm == "addr_3tau":
            return self.head(self.head_addr.signed_at(self.head_proj(h), self.taus))
        if self.arm == "addr_mhat":
            return self.head(self.head_addr.m_hat(self.head_proj(h)))
        if self.arm.startswith("addr_msl"):
            B, n, _ = h.shape
            slots = self.head_proj(h).view(B, n, self.n_slots, 4)
            if self.arm == "addr_msl_w":
                feats = self.head_addr.signed(slots).reshape(B, n, -1)
            elif getattr(self, "hard", False):
                feats = self.head_addr.m_hard_ste(slots).reshape(B, n, -1)
            else:
                feats = self.head_addr.m_hat(slots).reshape(B, n, -1)
            return self.head(feats)
        if self.arm == "addr_3tau_mhat":
            p = self.head_proj(h)
            feats = torch.cat([self.head_addr.signed_at(p, self.taus),
                               self.head_addr.m_hat(p)], dim=-1)
            return self.head(feats)
        return self.head(h)

    @torch.no_grad()
    def vitals(self) -> dict:
        out = {}
        if self.arm == "hub":
            for i, b in enumerate(self.blocks):
                if self._last_h is not None:
                    out[f"L{i}"] = b.attn.addr.vitals(b.attn.q(self._last_h[:2]))
        elif self.arm == "addr_head" and self._last_h is not None:
            out["head"] = self.head_addr.vitals(self._last_h[:2])
        elif self.arm in ("addr_d4", "addr_3tau", "addr_mhat",
                          "addr_3tau_mhat") and self._last_h is not None:
            out["head"] = self.head_addr.vitals(self.head_proj(self._last_h[:2]))
        elif self.arm.startswith("addr_msl") and self._last_h is not None:
            slots = self.head_proj(self._last_h[:2])
            out["head"] = self.head_addr.vitals(
                slots.reshape(*slots.shape[:-1], self.n_slots, 4))
        if self.use_relay and self._last_h is not None:
            for i, r in enumerate(self.relays):
                s = r.proj(self._last_h[:2])
                v = r.addr.vitals(s.reshape(*s.shape[:-1], r.n_slots, 4))
                out[f"relay{i}"] = {"gate": round(r.gate.sigmoid().item(), 4),
                                    "drift": v["drift"],
                                    "binding_frac": v["binding_frac"],
                                    "ppl": round(v["aliveness"]["usage_ppl"], 1)}
        return out


# --------------------------------------------------------------------------- data
def _wikitext_bytes(data_root: str):
    """wikitext-2-raw as flat uint8 tensors via the HF parquet CDN."""
    from huggingface_hub import hf_hub_download
    import pyarrow.parquet as pq

    def load(split):
        p = hf_hub_download("Salesforce/wikitext",
                            f"wikitext-2-raw-v1/{split}-00000-of-00001.parquet",
                            repo_type="dataset", local_dir=data_root)
        text = "".join(pq.read_table(p).column("text").to_pylist())
        return torch.frombuffer(bytearray(text.encode("utf-8")), dtype=torch.uint8).clone()

    return load("train"), load("validation")


def _batch(data: torch.Tensor, batch: int, block: int, device, g: torch.Generator):
    ix = torch.randint(0, data.numel() - block - 1, (batch,), generator=g)
    x = torch.stack([data[i:i + block] for i in ix]).long().to(device)
    y = torch.stack([data[i + 1:i + block + 1] for i in ix]).long().to(device)
    return x, y


# -------------------------------------------------------------------- train/smoke
def train(arms=("sdpa", "hub", "addr_head"), steps: int = 2000, batch: int = 32,
          block: int = 256, device: str = "cuda", data_root: str = "./data",
          seed: int = 0, eval_every: int = 500, save: bool = True):
    """Verdict run — GPU only. Pure Adam wd=0. Reports val bits-per-byte + vitals.
    save=True writes {data_root}/ar_ckpts/{arm}_s{seed}_t{steps}.pt per arm —
    the cultivated codebooks are SPECIMENS for the projective reading instruments."""
    import os
    if device == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("Verdict runs are GPU-only (never CPU-train for accuracy).")
    ckpt_dir = os.path.join(data_root, "ar_ckpts")
    os.makedirs(ckpt_dir, exist_ok=True)
    tr, va = _wikitext_bytes(data_root)
    print(f"data ready: train {tr.numel():,} bytes, val {va.numel():,} bytes", flush=True)
    results = {}
    for arm in arms:
        torch.manual_seed(seed)
        g = torch.Generator().manual_seed(seed)
        model = ByteLM(arm, block=block).to(device)
        n_params = sum(p.numel() for p in model.parameters())
        opt = torch.optim.Adam(model.parameters(), lr=3e-4, weight_decay=0.0)
        for step in range(1, steps + 1):
            x, y = _batch(tr, batch, block, device, g)
            logits = model(x)
            loss = F.cross_entropy(logits.reshape(-1, VOCAB), y.reshape(-1))
            opt.zero_grad(set_to_none=True)
            loss.backward()
            opt.step()
            if step % eval_every == 0 or step == steps:
                model.eval()
                with torch.no_grad():
                    losses = []
                    for _ in range(20):
                        xv, yv = _batch(va, batch, block, device, g)
                        lv = F.cross_entropy(model(xv).reshape(-1, VOCAB),
                                             yv.reshape(-1))
                        losses.append(lv.item())
                bpb = sum(losses) / len(losses) / math.log(2)
                print(f"[{arm}] step {step} val_bpb={bpb:.4f} vitals={model.vitals()}",
                      flush=True)
                model.train()
        results[arm] = {"val_bpb": bpb, "params": n_params, "vitals": model.vitals()}
        if save:
            path = os.path.join(ckpt_dir, f"{arm}_s{seed}_t{steps}.pt")
            torch.save({"arm": arm, "seed": seed, "steps": steps, "val_bpb": bpb,
                        "state_dict": {k: v.cpu() for k, v in
                                       model.state_dict().items()}}, path)
            print(f"saved specimen: {path}", flush=True)
    print(results, flush=True)
    return results


def smoke():
    """Shapes/parse only — no accuracy claims."""
    x = torch.randint(0, VOCAB, (2, 64))
    for arm in ("sdpa", "hub", "addr_head"):
        m = ByteLM(arm, d=96, layers=2, block=64, K=16)
        logits = m(x)
        assert logits.shape == (2, 64, VOCAB)
        logits.sum().backward()
        # causality check: future byte must not affect past logits
        with torch.no_grad():
            a = m(x)[0, 10]
            x2 = x.clone(); x2[0, 40] = (x2[0, 40] + 7) % 256
            b = m(x2)[0, 10]
        assert torch.allclose(a, b, atol=1e-4), f"{arm} leaks future context"
        print(f"{arm}: OK params={sum(p.numel() for p in m.parameters()):,} "
              f"vitals={m.vitals()}", flush=True)
    print("OK — AR bed smoke passed (verdict run: train() on GPU)", flush=True)


def _in_notebook() -> bool:
    try:
        get_ipython()  # type: ignore[name-defined]  # noqa: F821
        return True
    except NameError:
        return False


if __name__ == "__main__":
    if _in_notebook():
        smoke()
        print("Notebook mode: call  train(steps=2000)  in the next cell (GPU).")
    else:
        import argparse
        ap = argparse.ArgumentParser()
        ap.add_argument("--train", action="store_true")
        ap.add_argument("--steps", type=int, default=2000)
        a, _ = ap.parse_known_args()
        train(steps=a.steps) if a.train else smoke()