Prizma / seq /delta.py

Reframe as research artifact: rich card + Apache-2.0 license + clean runnable code subset; remove internal design files

91d91c5 verified 5 months ago

23.3 kB

"""
The Prizma-Seq workspace recurrence: a precision-gated delta rule (targeted erase-and-write),
derived as ONE gradient step on Prizma's per-token free energy  F_t(S) = 1/2 ||v_t - S k_t||^2.
    S_t = alpha_t * S_{t-1} + u_t k_t^T ,   u_t = beta_t (v_t - alpha_t S_{t-1} k_t)   (||k_t||=1)
    read (PRE-write, strictly causal):  o_t = S_{t-1} q_t
This module provides TWO implementations that MUST agree to < 1e-4 (forward AND grad):
  * _delta_reference  : naive per-token recurrence (ground truth, slow)
  * chunked_delta     : WY/UT chunk-parallel form (fast, used for training)
Both operate on [B,H,T,d_h] tensors. alpha defaults to 1 (pure DeltaNet) — the diagnostic gates
(MQAR/induction/selective-copy) need clean overwrite, not forgetting; the gated path is enabled
for char-LM. Keep d_h even, C<=64, float32.
Surprise-gated write (Lever A):
  When surprise=True, the write magnitude is modulated by the surprise gate
      g_t = 1 + tanh(||eps_t||)   in [1, 2)
      u_t = beta_t * g_t * eps_t
  derived from the per-token free energy F_t = ½||v_t - S k_t||^2.  Tokens with large prediction
  error write MORE (Titans-style test-time learning, predictive-coding grounded).
  SPEED NOTE: when surprise=True, chunked_delta internally runs an EXACT sequential scan (the WY/UT
  chunk-parallel shortcut cannot be used because eps_t = v_t - alpha*S_{t-1}*k_t depends on the
  FULL running state S_{t-1}, which is affected by every prior gated write — a frozen-chunk
  approximation diverges on REPEATED KEYS by ~100%).  This is a disclosed Pareto knob: correctness
  takes priority; a fused CUDA kernel can recover the speed later.  When surprise=False (default),
  the original fast WY/UT path is used byte-for-byte.
  surprise_mode controls HOW g_t is computed (R9 ablation controls):
    'norm'     : g_t = 1 + tanh(||eps_t||)                    — the real lever
    'random'   : g_t = 1 + tanh(|r_t|), r_t ~ |N(0,1)| via an explicit torch.Generator
                 (surprise_gen must be passed; reproducible).  Mean matches 'norm' heuristically:
                 E[tanh(|N(0,1)|)] ≈ 0.56 ≈ E[tanh(||eps||)] for typical normalised delta errors.
    'constant' : g_t = 1 + tanh(1.0) ≈ 1.762, a fixed constant equal to E[tanh(||eps||)]
                 approximated by tanh(1.0).  Provides a constant-mean-beta_eff control.
"""
from __future__ import annotations

import math
import torch


def _surprise_gate(eps, mode, gen):
    """Compute per-token surprise gate g = 1 + tanh(signal) in [1,2).
    eps: [B,H,d] prediction error. Returns g: [B,H,1] for broadcasting."""
    if mode == 'norm':
        # Real lever: g = 1 + tanh(||eps||);  ||eps|| is the free-energy sqrt.
        signal = eps.norm(dim=-1)                 # [B,H]
    elif mode == 'random':
        # Random-scalar control: g = 1 + tanh(|r|), r ~ |N(0,1)| via explicit generator.
        # E[tanh(|N(0,1)|)] ≈ 0.56, closely matching E[tanh(||eps||)] for typical delta errors
        # (both derived from half-normal; the match is heuristic — mean equality is approximate).
        assert gen is not None, "surprise_mode='random' requires surprise_gen"
        r = torch.zeros(eps.shape[:2], dtype=eps.dtype, device=eps.device)
        r.normal_(generator=gen)
        signal = r.abs()                          # [B,H], matched to |N(0,1)|
    elif mode == 'constant':
        # Constant-mean control: g = 1 + tanh(1.0) ≈ 1.762 for all tokens.
        # tanh(1.0) ≈ 0.762 approximates E[tanh(||eps||)] for a unit-scale error distribution.
        _CONST = math.tanh(1.0)
        return torch.full(eps.shape[:2] + (1,), 1.0 + _CONST, dtype=eps.dtype, device=eps.device)
    else:
        raise ValueError(f"Unknown surprise_mode: {mode!r}. Choose 'norm', 'random', or 'constant'.")
    return (1.0 + torch.tanh(signal))[..., None]  # [B,H,1]


def _delta_reference(q, k, v, beta, alpha=None, S0=None, write_mode="delta", beta_e=None,
                     n_delta=1, surprise=False, surprise_mode='norm', surprise_gen=None):
    """Ground-truth sequential recurrence. q,k,v:[B,H,T,d]; beta:[B,H,T]; alpha:[B,H,T] or None.
    write_mode='additive' -> u=beta*v (no erase), the linear-attn ablation.
    beta_e: optional erase gate [B,H,T]; if None, beta_e=beta (byte-identical to old behaviour).
    When beta_e is provided, u_t = beta_w*v_t - beta_e*(alpha*S k_t)  (decoupled GDN-2 write).
    n_delta: int >= 1 (DeltaProduct). For n_delta=1, inputs have no extra k-axis (byte-identical to
    old behaviour). For n_delta>=2, k,v have shape [B,H,T,n_delta,d] and beta [B,H,T,n_delta]:
    each token applies n_delta sequential delta sub-steps before reading the next token. The read
    o_t = S_{t-1} q_t uses the state from the END of the previous token's n_delta sub-steps. Alpha
    decay is applied exactly once per token (on sub-step 0 only); subsequent sub-steps are pure
    additive delta (no additional decay), consistent with Householder-product semantics.
    surprise: if True, scale u_t by g_t = 1+tanh(||eps_t||) (or per surprise_mode) BEFORE writing.
    surprise_mode: 'norm' | 'random' | 'constant' — see module docstring.
    surprise_gen: explicit torch.Generator for 'random' mode (reproducibility, R8 discipline)."""
    B, H, T, d = q.shape
    dv = v.shape[-1] if n_delta == 1 else v.shape[-1]  # [B,H,T,n_delta,d] -> last dim still d
    if n_delta >= 2:
        # v shape is [B,H,T,n_delta,d]; dv = d (last dim)
        dv = v.shape[-1]
    else:
        dv = v.shape[-1]
    if S0 is None:
        S = torch.zeros(B, H, dv, d, dtype=q.dtype, device=q.device)
    else:
        S = S0.clone()
    if alpha is None:
        alpha = torch.ones(B, H, T, dtype=q.dtype, device=q.device)
    if beta_e is None:
        beta_e = beta          # default: erase == write gate (byte-identical to today)
    erase = (write_mode == "delta")
    outs = []
    for t in range(T):
        qt = q[:, :, t]                                            # [B,H,d]
        at = alpha[:, :, t]                                        # [B,H]
        o = torch.einsum("bhij,bhj->bhi", S, qt)                   # read S_{t-1} q_t  (PRE-write)
        if n_delta == 1:
            kt, vt = k[:, :, t], v[:, :, t]                       # [B,H,d]
            bt  = beta[:, :, t]                                    # [B,H]
            bet = beta_e[:, :, t]
            if erase:
                Sk = torch.einsum("bhij,bhj->bhi", S, kt)          # [B,H,d]
                # Standard delta-rule error (free-energy gradient at S_{t-1}):
                #   eps_t = v_t - alpha_t * S_{t-1} k_t
                # Surprise gate g_t modulates ONLY the write magnitude (task spec: alpha unchanged):
                #   u_t = beta_t * g_t * eps_t  (when bet==bt, the default non-decoupled case)
                # For the decoupled-gate case (bet != bt, Lever H):
                #   the surprise gate is applied to the full u_t = bt*vt - bet*(at*Sk),
                #   which is the write vector. The eps for SIGNAL computation always uses
                #   the symmetric beta (bt), consistent with the free-energy interpretation.
                eps_t = vt - at[..., None] * Sk                    # [B,H,d] prediction error
                if surprise:
                    g = _surprise_gate(eps_t, surprise_mode, surprise_gen)  # [B,H,1]
                    # Apply g to the full write vector u (before alpha-scaled state decay)
                    u = g * (bt[..., None] * vt - bet[..., None] * (at[..., None] * Sk))
                else:
                    u = bt[..., None] * vt - bet[..., None] * (at[..., None] * Sk)
            else:
                u = bt[..., None] * vt
            S = at[..., None, None] * S + torch.einsum("bhi,bhj->bhij", u, kt)
        else:
            # n_delta >= 2: k[B,H,T,n_delta,d], v[B,H,T,n_delta,d], beta[B,H,T,n_delta]
            for j in range(n_delta):
                ktj = k[:, :, t, j]                               # [B,H,d]
                vtj = v[:, :, t, j]
                btj = beta[:, :, t, j]                             # [B,H]
                if j == 0:
                    # alpha decay applied once at the first sub-step
                    if erase:
                        Sk = torch.einsum("bhij,bhj->bhi", S, ktj)
                        u = btj[..., None] * vtj - btj[..., None] * (at[..., None] * Sk)
                    else:
                        u = btj[..., None] * vtj
                    S = at[..., None, None] * S + torch.einsum("bhi,bhj->bhij", u, ktj)
                else:
                    # subsequent sub-steps: pure delta, no additional alpha decay
                    if erase:
                        Sk = torch.einsum("bhij,bhj->bhi", S, ktj)
                        u = btj[..., None] * (vtj - Sk)
                    else:
                        u = btj[..., None] * vtj
                    S = S + torch.einsum("bhi,bhj->bhij", u, ktj)
        outs.append(o)
    O = torch.stack(outs, dim=2)                                   # [B,H,T,d]
    return O, S


def _solve_unit_lower(Amat, RHS):
    """Solve (I + A) X = RHS for X, where A is strictly-lower-triangular -> (I+A) unit-lower-tri.
    Tries solve_triangular (fast); falls back to an exact nilpotent Neumann series if MPS lacks it
    (A is strictly lower so A^C = 0; the series terminates and is exact)."""
    C = Amat.shape[-1]
    M = torch.eye(C, dtype=Amat.dtype, device=Amat.device) + Amat
    try:
        return torch.linalg.solve_triangular(M, RHS, upper=False, unitriangular=True)
    except Exception:
        # exact: (I+A)^{-1} = sum_{j=0}^{C-1} (-A)^j ; apply to RHS iteratively
        X = RHS.clone()
        term = RHS
        negA = -Amat
        for _ in range(1, C):
            term = negA @ term
            X = X + term
            if torch.count_nonzero(term) == 0:
                break
        return X


def chunked_delta(q, k, v, beta, alpha=None, S0=None, chunk=64, write_mode="delta", beta_e=None,
                  n_delta=1, surprise=False, surprise_mode='norm', surprise_gen=None):
    """WY/UT chunk-parallel delta rule. Same semantics as _delta_reference, O(T d^2/C + T C d).
    alpha=None -> pure delta (no decay). write_mode='additive' -> linear-attn ablation (no erase:
    u=beta*v, used for B6 PRIZMA_noDelta). Returns O:[B,H,T,d], S_end:[B,H,d,d].
    beta_e: optional erase gate [B,H,T]; if None, beta_e=beta (byte-identical to old behaviour).
    When beta_e is provided, the A-matrix (erase/read-back) is scaled by beta_e while the Vc
    (write) term is scaled by beta_w=beta — decoupled GDN-2 update.
    n_delta: int >= 1 (DeltaProduct). For n_delta=1, the standard WY/UT chunk-parallel path is
    used (byte-identical to old behaviour). For n_delta>=2, k,v have shape [B,H,T,n_delta,d] and
    beta [B,H,T,n_delta]. The k-axis is processed sequentially within each chunk (correctness-first
    fallback as documented in the plan) while T remains chunked. Alpha applied once per token.
    surprise: if True, delegate to _delta_reference for an EXACT sequential scan — the WY/UT
    chunk-parallel shortcut is NOT used because eps_t depends on S_{t-1}, which is affected by
    all prior surprise-gated writes. A frozen chunk-entry state approximation diverges on repeated
    keys by ~100% (remedy R3). When surprise=False (default), the fast WY/UT path is used
    byte-for-byte (no speed regression on the default path).
    surprise_mode / surprise_gen: forwarded to _delta_reference (see module docstring)."""
    B, H, T, d = q.shape
    dv = v.shape[-1]                     # value-dim-aware init -> RECTANGULAR state S in R^{d_v x d_k}
    if S0 is None:                       #   (d_k=d). Byte-identical when d_v == d_k (every existing
        S = torch.zeros(B, H, dv, d, dtype=q.dtype, device=q.device)   # call); enables feature-map/GDM
    else:
        S = S0

    # SURPRISE PATH: must be exact — delegate to _delta_reference which threads the TRUE running
    # state through every token.  The WY/UT chunk shortcut cannot be used here because eps_t depends
    # on S_{t-1}, which is modified by each surprise-gated write; a frozen chunk-entry state would
    # mis-estimate the surprise on repeated keys by ~100% (R3).  This is a disclosed speed cost;
    # a fused kernel can recover it later.  When surprise=False the fast path below is untouched.
    if surprise:
        return _delta_reference(q, k, v, beta, alpha=alpha, S0=S, write_mode=write_mode,
                                beta_e=beta_e, n_delta=n_delta,
                                surprise=True, surprise_mode=surprise_mode,
                                surprise_gen=surprise_gen)

    gated = alpha is not None
    erase = (write_mode == "delta")
    # beta_e=None means use beta for both erase and write -> byte-identical to previous behaviour
    decoupled = (beta_e is not None)

    # --- DeltaProduct n_delta >= 2: k-axis processed sequentially within each chunk ---
    # For n_delta=1 (default), fall through to the fast WY/UT path (byte-identical to old code).
    # For n_delta>=2: k,v are [B,H,T,n_delta,d], beta is [B,H,T,n_delta]. We chunk over T but
    # loop over tokens and sub-steps within each chunk (sequential k-axis, documented fallback).
    if n_delta >= 2:
        outs_nd = []
        for c0 in range(0, T, chunk):
            c1 = min(c0 + chunk, T)
            for ti in range(c0, c1):
                qt = q[:, :, ti]                                   # [B,H,d]
                at = alpha[:, :, ti] if gated else torch.ones(B, H, dtype=q.dtype, device=q.device)
                o = torch.einsum("bhij,bhj->bhi", S, qt)           # pre-write read
                for j in range(n_delta):
                    ktj = k[:, :, ti, j]                           # [B,H,d]
                    vtj = v[:, :, ti, j]
                    btj = beta[:, :, ti, j]                        # [B,H]
                    if j == 0:
                        # alpha applied once at first sub-step
                        if erase:
                            Sk = torch.einsum("bhij,bhj->bhi", S, ktj)
                            u = btj[..., None] * vtj - btj[..., None] * (at[..., None] * Sk)
                        else:
                            u = btj[..., None] * vtj
                        S = at[..., None, None] * S + torch.einsum("bhi,bhj->bhij", u, ktj)
                    else:
                        # subsequent sub-steps: pure delta, no alpha decay
                        if erase:
                            Sk = torch.einsum("bhij,bhj->bhi", S, ktj)
                            u = btj[..., None] * (vtj - Sk)
                        else:
                            u = btj[..., None] * vtj
                        S = S + torch.einsum("bhi,bhj->bhij", u, ktj)
                outs_nd.append(o)
        return torch.stack(outs_nd, dim=2), S

    outs = []
    for c0 in range(0, T, chunk):
        c1 = min(c0 + chunk, T)
        C = c1 - c0
        Kc = k[:, :, c0:c1]                      # [B,H,C,d]
        Vc = v[:, :, c0:c1]
        Qc = q[:, :, c0:c1]
        Bc = beta[:, :, c0:c1]                    # [B,H,C]  write gate beta_w
        Bec = beta_e[:, :, c0:c1] if decoupled else Bc   # [B,H,C]  erase gate beta_e
        if gated:
            Ac = alpha[:, :, c0:c1]               # [B,H,C]  in [0.5,1]
            # within-chunk cumulative decay in LOG space (avoids float32 underflow of gamma over
            # long chunks: exp(cumsum(log a)) -> 1e-20 at C=64, and dividing two underflowed
            # exponentials destroys the recurrence). All gamma-RATIOS are formed as exp(log-diff),
            # which is <=1 on the strictly-lower region actually used -> stable, no clamp.
            logA = torch.log(Ac.clamp_min(1e-6))
            clog = torch.cumsum(logA, dim=-1)                    # [B,H,C]  log gamma_i (post i)
            clog_prev = clog - logA                              # log gamma_{i-1} (pre i)
            KK = torch.matmul(Kc, Kc.transpose(-1, -2))          # [B,H,C,C]  k_i·k_j
            ratio = torch.exp(clog[..., :, None] - clog[..., None, :])         # gamma_i/gamma_j
            # A matrix: erase gate beta_e scales the KK*ratio terms (read-back / erase strength)
            A = torch.tril(Bec[..., :, None] * (KK * ratio), -1) if erase else torch.zeros_like(KK)
            KS0 = torch.matmul(Kc, S.transpose(-1, -2))          # [B,H,C,d]  (k_i^T S0^T)
            gamma = torch.exp(clog)[..., None]                   # [B,H,C,1] absolute (genuine small)
            # rhs: write gate beta_w scales Vc; erase gate beta_e scales the gamma*KS0 term
            if erase:
                rhs = Bc[..., None] * Vc - Bec[..., None] * (gamma * KS0)
            else:
                rhs = Bc[..., None] * Vc
            U = _solve_unit_lower(A, rhs)                        # [B,H,C,d]
            # reads are PRE-write -> decayed to gamma_{i-1}
            read_ratio = torch.exp(clog_prev[..., :, None] - clog[..., None, :])   # gamma_{i-1}/gamma_j
            O_inter = torch.exp(clog_prev)[..., None] * torch.matmul(Qc, S.transpose(-1, -2))
            QK = torch.matmul(Qc, Kc.transpose(-1, -2)) * read_ratio
            O_intra = torch.matmul(torch.tril(QK, -1), U)
            Oc = O_inter + O_intra
            # state carry: S_end = gamma_C S0 + sum_i (gamma_C/gamma_i) u_i k_i^T (ratio <=1 -> stable)
            clogC = clog[..., -1:]                               # [B,H,1]
            gC = torch.exp(clogC)[..., None]                     # [B,H,1,1]
            scale = torch.exp(clogC - clog)                      # [B,H,C]  gamma_C/gamma_i <= 1
            S = gC * S + torch.matmul((scale[..., None] * U).transpose(-1, -2), Kc)
        else:
            KK = torch.matmul(Kc, Kc.transpose(-1, -2))          # [B,H,C,C]
            # A matrix: erase gate beta_e scales the KK terms
            A = torch.tril(Bec[..., :, None] * KK, -1) if erase else torch.zeros_like(KK)
            KS0 = torch.matmul(Kc, S.transpose(-1, -2))          # [B,H,C,d]
            # rhs: write gate beta_w scales Vc; erase gate beta_e scales KS0
            if erase:
                rhs = Bc[..., None] * Vc - Bec[..., None] * KS0
            else:
                rhs = Bc[..., None] * Vc
            U = _solve_unit_lower(A, rhs)                        # [B,H,C,d]
            O_inter = torch.matmul(Qc, S.transpose(-1, -2))      # [B,H,C,d]
            QK = torch.matmul(Qc, Kc.transpose(-1, -2))
            O_intra = torch.matmul(torch.tril(QK, -1), U)
            Oc = O_inter + O_intra
            S = S + torch.matmul(U.transpose(-1, -2), Kc)        # S0 + U^T K
        outs.append(Oc)
    return torch.cat(outs, dim=2), S


if __name__ == "__main__":
    torch.manual_seed(0)
    # cover the PRODUCTION regime: chunk=64, long T, the alpha=0.5 decay floor (worst-case
    # underflow), and the additive ablation. (The old test used chunk=16/T=40 and HID a float32
    # underflow bug in the gated chunked path — never weaken these settings.)
    devs = ["cpu", "mps"] if torch.backends.mps.is_available() else ["cpu"]
    cases = [("pure", None, "delta"), ("gated~U", "rand", "delta"),
             ("gated.5floor", "floor", "delta"), ("additive", None, "additive"),
             ("gated-additive", "rand", "additive")]
    allok = True
    for name, amode, wmode in cases:
        for dev in devs:
            B, H, T, d, C = 2, 3, 256, 16, 64
            q = torch.randn(B, H, T, d, device=dev)
            k = torch.randn(B, H, T, d, device=dev); k = k / k.norm(dim=-1, keepdim=True)
            v = torch.randn(B, H, T, d, device=dev)
            beta = torch.rand(B, H, T, device=dev) * 0.99
            if amode == "rand":
                alpha = 0.5 + 0.5 * torch.rand(B, H, T, device=dev)
            elif amode == "floor":
                alpha = torch.full((B, H, T), 0.5, device=dev)     # worst-case sustained decay
            else:
                alpha = None
            Oref, Sref = _delta_reference(q, k, v, beta, alpha, write_mode=wmode)
            Och, Sch = chunked_delta(q, k, v, beta, alpha, chunk=C, write_mode=wmode)
            do = (Oref - Och).abs().max().item(); ds = (Sref - Sch).abs().max().item()
            ok = max(do, ds) < 1e-4; allok &= ok
            print(f"[{name:<14} {dev}] C={C} T={T} max|dO|={do:.2e} max|dS|={ds:.2e}  "
                  f"{'OK' if ok else 'MISMATCH'}")
    # RECTANGULAR state (d_v != d_k): the feature-map / GlobalDeltaMemory contract. The chunked
    # WY/UT form must equal the naive recurrence when keys/queries have a different last-dim than
    # values (state S in R^{d_v x d_k}). This is the load-bearing guarantee for the capacity levers.
    for dev in devs:
        B, H, T, dk, dv, C = 2, 3, 200, 48, 16, 64
        q = torch.randn(B, H, T, dk, device=dev)
        k = torch.randn(B, H, T, dk, device=dev); k = k / k.norm(dim=-1, keepdim=True)
        v = torch.randn(B, H, T, dv, device=dev)
        beta = torch.rand(B, H, T, device=dev) * 0.99
        for amode, wmode in [(None, "delta"), ("rand", "delta"), (None, "additive")]:
            alpha = (0.5 + 0.5 * torch.rand(B, H, T, device=dev)) if amode == "rand" else None
            Oref, Sref = _delta_reference(q, k, v, beta, alpha, write_mode=wmode)
            Och, Sch = chunked_delta(q, k, v, beta, alpha, chunk=C, write_mode=wmode)
            do = (Oref - Och).abs().max().item(); ds = (Sref - Sch).abs().max().item()
            ok = max(do, ds) < 1e-4; allok &= ok
            print(f"[rect dv={dv}!=dk={dk} {dev} {wmode}{'/gated' if amode else ''}] "
                  f"max|dO|={do:.2e} max|dS|={ds:.2e}  {'OK' if ok else 'MISMATCH'}")
    # SURPRISE GATE: check that chunked_delta(surprise=True) == _delta_reference(surprise=True)
    # including on REPEATED KEYS (the R3 binding test).
    for dev in devs:
        B, H, T, d, C = 2, 2, 64, 16, 32
        torch.manual_seed(99)
        q = torch.randn(B, H, T, d, device=dev)
        # Repeat a fixed key for the first half (worst-case R3 scenario):
        k_rep = torch.randn(1, 1, 1, d, device=dev)
        k_rep = k_rep / k_rep.norm(dim=-1, keepdim=True)
        k = k_rep.expand(B, H, T, d).clone()   # all tokens share the SAME key
        v = torch.randn(B, H, T, d, device=dev)
        beta = torch.rand(B, H, T, device=dev) * 0.99
        alpha_surp = 0.5 + 0.5 * torch.rand(B, H, T, device=dev)
        for smode in ('norm', 'constant'):
            gen = torch.Generator().manual_seed(42) if smode == 'random' else None
            Oref_s, Sref_s = _delta_reference(q, k, v, beta, alpha_surp,
                                               surprise=True, surprise_mode=smode, surprise_gen=gen)
            gen2 = torch.Generator().manual_seed(42) if smode == 'random' else None
            Och_s, Sch_s = chunked_delta(q, k, v, beta, alpha_surp, chunk=C,
                                          surprise=True, surprise_mode=smode, surprise_gen=gen2)
            do_s = (Oref_s - Och_s).abs().max().item()
            ds_s = (Sref_s - Sch_s).abs().max().item()
            ok_s = max(do_s, ds_s) < 1e-4; allok &= ok_s
            print(f"[surprise={smode:<8} {dev}] repeated-key max|dO|={do_s:.2e} max|dS|={ds_s:.2e}  "
                  f"{'OK' if ok_s else 'MISMATCH'}")
    print("ALL OK" if allok else "FAILURES PRESENT")

	"""
	The Prizma-Seq workspace recurrence: a precision-gated delta rule (targeted erase-and-write),
	derived as ONE gradient step on Prizma's per-token free energy F_t(S) = 1/2 \|\|v_t - S k_t\|\|^2.

	S_t = alpha_t * S_{t-1} + u_t k_t^T , u_t = beta_t (v_t - alpha_t S_{t-1} k_t) (\|\|k_t\|\|=1)
	read (PRE-write, strictly causal): o_t = S_{t-1} q_t

	This module provides TWO implementations that MUST agree to < 1e-4 (forward AND grad):
	* _delta_reference : naive per-token recurrence (ground truth, slow)
	* chunked_delta : WY/UT chunk-parallel form (fast, used for training)
	Both operate on [B,H,T,d_h] tensors. alpha defaults to 1 (pure DeltaNet) — the diagnostic gates
	(MQAR/induction/selective-copy) need clean overwrite, not forgetting; the gated path is enabled
	for char-LM. Keep d_h even, C<=64, float32.

	Surprise-gated write (Lever A):
	When surprise=True, the write magnitude is modulated by the surprise gate
	g_t = 1 + tanh(\|\|eps_t\|\|) in [1, 2)
	u_t = beta_t * g_t * eps_t
	derived from the per-token free energy F_t = ½\|\|v_t - S k_t\|\|^2. Tokens with large prediction
	error write MORE (Titans-style test-time learning, predictive-coding grounded).

	SPEED NOTE: when surprise=True, chunked_delta internally runs an EXACT sequential scan (the WY/UT
	chunk-parallel shortcut cannot be used because eps_t = v_t - alphaS_{t-1}k_t depends on the
	FULL running state S_{t-1}, which is affected by every prior gated write — a frozen-chunk
	approximation diverges on REPEATED KEYS by ~100%). This is a disclosed Pareto knob: correctness
	takes priority; a fused CUDA kernel can recover the speed later. When surprise=False (default),
	the original fast WY/UT path is used byte-for-byte.

	surprise_mode controls HOW g_t is computed (R9 ablation controls):
	'norm' : g_t = 1 + tanh(\|\|eps_t\|\|) — the real lever
	'random' : g_t = 1 + tanh(\|r_t\|), r_t ~ \|N(0,1)\| via an explicit torch.Generator
	(surprise_gen must be passed; reproducible). Mean matches 'norm' heuristically:
	E[tanh(\|N(0,1)\|)] ≈ 0.56 ≈ E[tanh(\|\|eps\|\|)] for typical normalised delta errors.
	'constant' : g_t = 1 + tanh(1.0) ≈ 1.762, a fixed constant equal to E[tanh(\|\|eps\|\|)]
	approximated by tanh(1.0). Provides a constant-mean-beta_eff control.
	"""
	from __future__ import annotations

	import math
	import torch


	def _surprise_gate(eps, mode, gen):
	"""Compute per-token surprise gate g = 1 + tanh(signal) in [1,2).
	eps: [B,H,d] prediction error. Returns g: [B,H,1] for broadcasting."""
	if mode == 'norm':
	# Real lever: g = 1 + tanh(\|\|eps\|\|); \|\|eps\|\| is the free-energy sqrt.
	signal = eps.norm(dim=-1) # [B,H]
	elif mode == 'random':
	# Random-scalar control: g = 1 + tanh(\|r\|), r ~ \|N(0,1)\| via explicit generator.
	# E[tanh(\|N(0,1)\|)] ≈ 0.56, closely matching E[tanh(\|\|eps\|\|)] for typical delta errors
	# (both derived from half-normal; the match is heuristic — mean equality is approximate).
	assert gen is not None, "surprise_mode='random' requires surprise_gen"
	r = torch.zeros(eps.shape[:2], dtype=eps.dtype, device=eps.device)
	r.normal_(generator=gen)
	signal = r.abs() # [B,H], matched to \|N(0,1)\|
	elif mode == 'constant':
	# Constant-mean control: g = 1 + tanh(1.0) ≈ 1.762 for all tokens.
	# tanh(1.0) ≈ 0.762 approximates E[tanh(\|\|eps\|\|)] for a unit-scale error distribution.
	_CONST = math.tanh(1.0)
	return torch.full(eps.shape[:2] + (1,), 1.0 + _CONST, dtype=eps.dtype, device=eps.device)
	else:
	raise ValueError(f"Unknown surprise_mode: {mode!r}. Choose 'norm', 'random', or 'constant'.")
	return (1.0 + torch.tanh(signal))[..., None] # [B,H,1]


	def _delta_reference(q, k, v, beta, alpha=None, S0=None, write_mode="delta", beta_e=None,
	n_delta=1, surprise=False, surprise_mode='norm', surprise_gen=None):
	"""Ground-truth sequential recurrence. q,k,v:[B,H,T,d]; beta:[B,H,T]; alpha:[B,H,T] or None.
	write_mode='additive' -> u=beta*v (no erase), the linear-attn ablation.
	beta_e: optional erase gate [B,H,T]; if None, beta_e=beta (byte-identical to old behaviour).
	When beta_e is provided, u_t = beta_wv_t - beta_e(alpha*S k_t) (decoupled GDN-2 write).
	n_delta: int >= 1 (DeltaProduct). For n_delta=1, inputs have no extra k-axis (byte-identical to
	old behaviour). For n_delta>=2, k,v have shape [B,H,T,n_delta,d] and beta [B,H,T,n_delta]:
	each token applies n_delta sequential delta sub-steps before reading the next token. The read
	o_t = S_{t-1} q_t uses the state from the END of the previous token's n_delta sub-steps. Alpha
	decay is applied exactly once per token (on sub-step 0 only); subsequent sub-steps are pure
	additive delta (no additional decay), consistent with Householder-product semantics.
	surprise: if True, scale u_t by g_t = 1+tanh(\|\|eps_t\|\|) (or per surprise_mode) BEFORE writing.
	surprise_mode: 'norm' \| 'random' \| 'constant' — see module docstring.
	surprise_gen: explicit torch.Generator for 'random' mode (reproducibility, R8 discipline)."""
	B, H, T, d = q.shape
	dv = v.shape[-1] if n_delta == 1 else v.shape[-1] # [B,H,T,n_delta,d] -> last dim still d
	if n_delta >= 2:
	# v shape is [B,H,T,n_delta,d]; dv = d (last dim)
	dv = v.shape[-1]
	else:
	dv = v.shape[-1]
	if S0 is None:
	S = torch.zeros(B, H, dv, d, dtype=q.dtype, device=q.device)
	else:
	S = S0.clone()
	if alpha is None:
	alpha = torch.ones(B, H, T, dtype=q.dtype, device=q.device)
	if beta_e is None:
	beta_e = beta # default: erase == write gate (byte-identical to today)
	erase = (write_mode == "delta")
	outs = []
	for t in range(T):
	qt = q[:, :, t] # [B,H,d]
	at = alpha[:, :, t] # [B,H]
	o = torch.einsum("bhij,bhj->bhi", S, qt) # read S_{t-1} q_t (PRE-write)
	if n_delta == 1:
	kt, vt = k[:, :, t], v[:, :, t] # [B,H,d]
	bt = beta[:, :, t] # [B,H]
	bet = beta_e[:, :, t]
	if erase:
	Sk = torch.einsum("bhij,bhj->bhi", S, kt) # [B,H,d]
	# Standard delta-rule error (free-energy gradient at S_{t-1}):
	# eps_t = v_t - alpha_t * S_{t-1} k_t
	# Surprise gate g_t modulates ONLY the write magnitude (task spec: alpha unchanged):
	# u_t = beta_t * g_t * eps_t (when bet==bt, the default non-decoupled case)
	# For the decoupled-gate case (bet != bt, Lever H):
	# the surprise gate is applied to the full u_t = btvt - bet(at*Sk),
	# which is the write vector. The eps for SIGNAL computation always uses
	# the symmetric beta (bt), consistent with the free-energy interpretation.
	eps_t = vt - at[..., None] * Sk # [B,H,d] prediction error
	if surprise:
	g = _surprise_gate(eps_t, surprise_mode, surprise_gen) # [B,H,1]
	# Apply g to the full write vector u (before alpha-scaled state decay)
	u = g * (bt[..., None] * vt - bet[..., None] * (at[..., None] * Sk))
	else:
	u = bt[..., None] * vt - bet[..., None] * (at[..., None] * Sk)
	else:
	u = bt[..., None] * vt
	S = at[..., None, None] * S + torch.einsum("bhi,bhj->bhij", u, kt)
	else:
	# n_delta >= 2: k[B,H,T,n_delta,d], v[B,H,T,n_delta,d], beta[B,H,T,n_delta]
	for j in range(n_delta):
	ktj = k[:, :, t, j] # [B,H,d]
	vtj = v[:, :, t, j]
	btj = beta[:, :, t, j] # [B,H]
	if j == 0:
	# alpha decay applied once at the first sub-step
	if erase:
	Sk = torch.einsum("bhij,bhj->bhi", S, ktj)
	u = btj[..., None] * vtj - btj[..., None] * (at[..., None] * Sk)
	else:
	u = btj[..., None] * vtj
	S = at[..., None, None] * S + torch.einsum("bhi,bhj->bhij", u, ktj)
	else:
	# subsequent sub-steps: pure delta, no additional alpha decay
	if erase:
	Sk = torch.einsum("bhij,bhj->bhi", S, ktj)
	u = btj[..., None] * (vtj - Sk)
	else:
	u = btj[..., None] * vtj
	S = S + torch.einsum("bhi,bhj->bhij", u, ktj)
	outs.append(o)
	O = torch.stack(outs, dim=2) # [B,H,T,d]
	return O, S


	def _solve_unit_lower(Amat, RHS):
	"""Solve (I + A) X = RHS for X, where A is strictly-lower-triangular -> (I+A) unit-lower-tri.
	Tries solve_triangular (fast); falls back to an exact nilpotent Neumann series if MPS lacks it
	(A is strictly lower so A^C = 0; the series terminates and is exact)."""
	C = Amat.shape[-1]
	M = torch.eye(C, dtype=Amat.dtype, device=Amat.device) + Amat
	try:
	return torch.linalg.solve_triangular(M, RHS, upper=False, unitriangular=True)
	except Exception:
	# exact: (I+A)^{-1} = sum_{j=0}^{C-1} (-A)^j ; apply to RHS iteratively
	X = RHS.clone()
	term = RHS
	negA = -Amat
	for _ in range(1, C):
	term = negA @ term
	X = X + term
	if torch.count_nonzero(term) == 0:
	break
	return X


	def chunked_delta(q, k, v, beta, alpha=None, S0=None, chunk=64, write_mode="delta", beta_e=None,
	n_delta=1, surprise=False, surprise_mode='norm', surprise_gen=None):
	"""WY/UT chunk-parallel delta rule. Same semantics as _delta_reference, O(T d^2/C + T C d).
	alpha=None -> pure delta (no decay). write_mode='additive' -> linear-attn ablation (no erase:
	u=beta*v, used for B6 PRIZMA_noDelta). Returns O:[B,H,T,d], S_end:[B,H,d,d].
	beta_e: optional erase gate [B,H,T]; if None, beta_e=beta (byte-identical to old behaviour).
	When beta_e is provided, the A-matrix (erase/read-back) is scaled by beta_e while the Vc
	(write) term is scaled by beta_w=beta — decoupled GDN-2 update.
	n_delta: int >= 1 (DeltaProduct). For n_delta=1, the standard WY/UT chunk-parallel path is
	used (byte-identical to old behaviour). For n_delta>=2, k,v have shape [B,H,T,n_delta,d] and
	beta [B,H,T,n_delta]. The k-axis is processed sequentially within each chunk (correctness-first
	fallback as documented in the plan) while T remains chunked. Alpha applied once per token.
	surprise: if True, delegate to _delta_reference for an EXACT sequential scan — the WY/UT
	chunk-parallel shortcut is NOT used because eps_t depends on S_{t-1}, which is affected by
	all prior surprise-gated writes. A frozen chunk-entry state approximation diverges on repeated
	keys by ~100% (remedy R3). When surprise=False (default), the fast WY/UT path is used
	byte-for-byte (no speed regression on the default path).
	surprise_mode / surprise_gen: forwarded to _delta_reference (see module docstring)."""
	B, H, T, d = q.shape
	dv = v.shape[-1] # value-dim-aware init -> RECTANGULAR state S in R^{d_v x d_k}
	if S0 is None: # (d_k=d). Byte-identical when d_v == d_k (every existing
	S = torch.zeros(B, H, dv, d, dtype=q.dtype, device=q.device) # call); enables feature-map/GDM
	else:
	S = S0

	# SURPRISE PATH: must be exact — delegate to _delta_reference which threads the TRUE running
	# state through every token. The WY/UT chunk shortcut cannot be used here because eps_t depends
	# on S_{t-1}, which is modified by each surprise-gated write; a frozen chunk-entry state would
	# mis-estimate the surprise on repeated keys by ~100% (R3). This is a disclosed speed cost;
	# a fused kernel can recover it later. When surprise=False the fast path below is untouched.
	if surprise:
	return _delta_reference(q, k, v, beta, alpha=alpha, S0=S, write_mode=write_mode,
	beta_e=beta_e, n_delta=n_delta,
	surprise=True, surprise_mode=surprise_mode,
	surprise_gen=surprise_gen)

	gated = alpha is not None
	erase = (write_mode == "delta")
	# beta_e=None means use beta for both erase and write -> byte-identical to previous behaviour
	decoupled = (beta_e is not None)

	# --- DeltaProduct n_delta >= 2: k-axis processed sequentially within each chunk ---
	# For n_delta=1 (default), fall through to the fast WY/UT path (byte-identical to old code).
	# For n_delta>=2: k,v are [B,H,T,n_delta,d], beta is [B,H,T,n_delta]. We chunk over T but
	# loop over tokens and sub-steps within each chunk (sequential k-axis, documented fallback).
	if n_delta >= 2:
	outs_nd = []
	for c0 in range(0, T, chunk):
	c1 = min(c0 + chunk, T)
	for ti in range(c0, c1):
	qt = q[:, :, ti] # [B,H,d]
	at = alpha[:, :, ti] if gated else torch.ones(B, H, dtype=q.dtype, device=q.device)
	o = torch.einsum("bhij,bhj->bhi", S, qt) # pre-write read
	for j in range(n_delta):
	ktj = k[:, :, ti, j] # [B,H,d]
	vtj = v[:, :, ti, j]
	btj = beta[:, :, ti, j] # [B,H]
	if j == 0:
	# alpha applied once at first sub-step
	if erase:
	Sk = torch.einsum("bhij,bhj->bhi", S, ktj)
	u = btj[..., None] * vtj - btj[..., None] * (at[..., None] * Sk)
	else:
	u = btj[..., None] * vtj
	S = at[..., None, None] * S + torch.einsum("bhi,bhj->bhij", u, ktj)
	else:
	# subsequent sub-steps: pure delta, no alpha decay
	if erase:
	Sk = torch.einsum("bhij,bhj->bhi", S, ktj)
	u = btj[..., None] * (vtj - Sk)
	else:
	u = btj[..., None] * vtj
	S = S + torch.einsum("bhi,bhj->bhij", u, ktj)
	outs_nd.append(o)
	return torch.stack(outs_nd, dim=2), S

	outs = []
	for c0 in range(0, T, chunk):
	c1 = min(c0 + chunk, T)
	C = c1 - c0
	Kc = k[:, :, c0:c1] # [B,H,C,d]
	Vc = v[:, :, c0:c1]
	Qc = q[:, :, c0:c1]
	Bc = beta[:, :, c0:c1] # [B,H,C] write gate beta_w
	Bec = beta_e[:, :, c0:c1] if decoupled else Bc # [B,H,C] erase gate beta_e
	if gated:
	Ac = alpha[:, :, c0:c1] # [B,H,C] in [0.5,1]
	# within-chunk cumulative decay in LOG space (avoids float32 underflow of gamma over
	# long chunks: exp(cumsum(log a)) -> 1e-20 at C=64, and dividing two underflowed
	# exponentials destroys the recurrence). All gamma-RATIOS are formed as exp(log-diff),
	# which is <=1 on the strictly-lower region actually used -> stable, no clamp.
	logA = torch.log(Ac.clamp_min(1e-6))
	clog = torch.cumsum(logA, dim=-1) # [B,H,C] log gamma_i (post i)
	clog_prev = clog - logA # log gamma_{i-1} (pre i)
	KK = torch.matmul(Kc, Kc.transpose(-1, -2)) # [B,H,C,C] k_i·k_j
	ratio = torch.exp(clog[..., :, None] - clog[..., None, :]) # gamma_i/gamma_j
	# A matrix: erase gate beta_e scales the KK*ratio terms (read-back / erase strength)
	A = torch.tril(Bec[..., :, None] * (KK * ratio), -1) if erase else torch.zeros_like(KK)
	KS0 = torch.matmul(Kc, S.transpose(-1, -2)) # [B,H,C,d] (k_i^T S0^T)
	gamma = torch.exp(clog)[..., None] # [B,H,C,1] absolute (genuine small)
	# rhs: write gate beta_w scales Vc; erase gate beta_e scales the gamma*KS0 term
	if erase:
	rhs = Bc[..., None] * Vc - Bec[..., None] * (gamma * KS0)
	else:
	rhs = Bc[..., None] * Vc
	U = _solve_unit_lower(A, rhs) # [B,H,C,d]
	# reads are PRE-write -> decayed to gamma_{i-1}
	read_ratio = torch.exp(clog_prev[..., :, None] - clog[..., None, :]) # gamma_{i-1}/gamma_j
	O_inter = torch.exp(clog_prev)[..., None] * torch.matmul(Qc, S.transpose(-1, -2))
	QK = torch.matmul(Qc, Kc.transpose(-1, -2)) * read_ratio
	O_intra = torch.matmul(torch.tril(QK, -1), U)
	Oc = O_inter + O_intra
	# state carry: S_end = gamma_C S0 + sum_i (gamma_C/gamma_i) u_i k_i^T (ratio <=1 -> stable)
	clogC = clog[..., -1:] # [B,H,1]
	gC = torch.exp(clogC)[..., None] # [B,H,1,1]
	scale = torch.exp(clogC - clog) # [B,H,C] gamma_C/gamma_i <= 1
	S = gC * S + torch.matmul((scale[..., None] * U).transpose(-1, -2), Kc)
	else:
	KK = torch.matmul(Kc, Kc.transpose(-1, -2)) # [B,H,C,C]
	# A matrix: erase gate beta_e scales the KK terms
	A = torch.tril(Bec[..., :, None] * KK, -1) if erase else torch.zeros_like(KK)
	KS0 = torch.matmul(Kc, S.transpose(-1, -2)) # [B,H,C,d]
	# rhs: write gate beta_w scales Vc; erase gate beta_e scales KS0
	if erase:
	rhs = Bc[..., None] * Vc - Bec[..., None] * KS0
	else:
	rhs = Bc[..., None] * Vc
	U = _solve_unit_lower(A, rhs) # [B,H,C,d]
	O_inter = torch.matmul(Qc, S.transpose(-1, -2)) # [B,H,C,d]
	QK = torch.matmul(Qc, Kc.transpose(-1, -2))
	O_intra = torch.matmul(torch.tril(QK, -1), U)
	Oc = O_inter + O_intra
	S = S + torch.matmul(U.transpose(-1, -2), Kc) # S0 + U^T K
	outs.append(Oc)
	return torch.cat(outs, dim=2), S


	if __name__ == "__main__":
	torch.manual_seed(0)
	# cover the PRODUCTION regime: chunk=64, long T, the alpha=0.5 decay floor (worst-case
	# underflow), and the additive ablation. (The old test used chunk=16/T=40 and HID a float32
	# underflow bug in the gated chunked path — never weaken these settings.)
	devs = ["cpu", "mps"] if torch.backends.mps.is_available() else ["cpu"]
	cases = [("pure", None, "delta"), ("gated~U", "rand", "delta"),
	("gated.5floor", "floor", "delta"), ("additive", None, "additive"),
	("gated-additive", "rand", "additive")]
	allok = True
	for name, amode, wmode in cases:
	for dev in devs:
	B, H, T, d, C = 2, 3, 256, 16, 64
	q = torch.randn(B, H, T, d, device=dev)
	k = torch.randn(B, H, T, d, device=dev); k = k / k.norm(dim=-1, keepdim=True)
	v = torch.randn(B, H, T, d, device=dev)
	beta = torch.rand(B, H, T, device=dev) * 0.99
	if amode == "rand":
	alpha = 0.5 + 0.5 * torch.rand(B, H, T, device=dev)
	elif amode == "floor":
	alpha = torch.full((B, H, T), 0.5, device=dev) # worst-case sustained decay
	else:
	alpha = None
	Oref, Sref = _delta_reference(q, k, v, beta, alpha, write_mode=wmode)
	Och, Sch = chunked_delta(q, k, v, beta, alpha, chunk=C, write_mode=wmode)
	do = (Oref - Och).abs().max().item(); ds = (Sref - Sch).abs().max().item()
	ok = max(do, ds) < 1e-4; allok &= ok
	print(f"[{name:<14} {dev}] C={C} T={T} max\|dO\|={do:.2e} max\|dS\|={ds:.2e} "
	f"{'OK' if ok else 'MISMATCH'}")
	# RECTANGULAR state (d_v != d_k): the feature-map / GlobalDeltaMemory contract. The chunked
	# WY/UT form must equal the naive recurrence when keys/queries have a different last-dim than
	# values (state S in R^{d_v x d_k}). This is the load-bearing guarantee for the capacity levers.
	for dev in devs:
	B, H, T, dk, dv, C = 2, 3, 200, 48, 16, 64
	q = torch.randn(B, H, T, dk, device=dev)
	k = torch.randn(B, H, T, dk, device=dev); k = k / k.norm(dim=-1, keepdim=True)
	v = torch.randn(B, H, T, dv, device=dev)
	beta = torch.rand(B, H, T, device=dev) * 0.99
	for amode, wmode in [(None, "delta"), ("rand", "delta"), (None, "additive")]:
	alpha = (0.5 + 0.5 * torch.rand(B, H, T, device=dev)) if amode == "rand" else None
	Oref, Sref = _delta_reference(q, k, v, beta, alpha, write_mode=wmode)
	Och, Sch = chunked_delta(q, k, v, beta, alpha, chunk=C, write_mode=wmode)
	do = (Oref - Och).abs().max().item(); ds = (Sref - Sch).abs().max().item()
	ok = max(do, ds) < 1e-4; allok &= ok
	print(f"[rect dv={dv}!=dk={dk} {dev} {wmode}{'/gated' if amode else ''}] "
	f"max\|dO\|={do:.2e} max\|dS\|={ds:.2e} {'OK' if ok else 'MISMATCH'}")
	# SURPRISE GATE: check that chunked_delta(surprise=True) == _delta_reference(surprise=True)
	# including on REPEATED KEYS (the R3 binding test).
	for dev in devs:
	B, H, T, d, C = 2, 2, 64, 16, 32
	torch.manual_seed(99)
	q = torch.randn(B, H, T, d, device=dev)
	# Repeat a fixed key for the first half (worst-case R3 scenario):
	k_rep = torch.randn(1, 1, 1, d, device=dev)
	k_rep = k_rep / k_rep.norm(dim=-1, keepdim=True)
	k = k_rep.expand(B, H, T, d).clone() # all tokens share the SAME key
	v = torch.randn(B, H, T, d, device=dev)
	beta = torch.rand(B, H, T, device=dev) * 0.99
	alpha_surp = 0.5 + 0.5 * torch.rand(B, H, T, device=dev)
	for smode in ('norm', 'constant'):
	gen = torch.Generator().manual_seed(42) if smode == 'random' else None
	Oref_s, Sref_s = _delta_reference(q, k, v, beta, alpha_surp,
	surprise=True, surprise_mode=smode, surprise_gen=gen)
	gen2 = torch.Generator().manual_seed(42) if smode == 'random' else None
	Och_s, Sch_s = chunked_delta(q, k, v, beta, alpha_surp, chunk=C,
	surprise=True, surprise_mode=smode, surprise_gen=gen2)
	do_s = (Oref_s - Och_s).abs().max().item()
	ds_s = (Sref_s - Sch_s).abs().max().item()
	ok_s = max(do_s, ds_s) < 1e-4; allok &= ok_s
	print(f"[surprise={smode:<8} {dev}] repeated-key max\|dO\|={do_s:.2e} max\|dS\|={ds_s:.2e} "
	f"{'OK' if ok_s else 'MISMATCH'}")
	print("ALL OK" if allok else "FAILURES PRESENT")