1"""Mandatory invariant suite (WP31 brief, Phase 3). Each test maps to one numbered invariant."""2import copy34import pytest5import torch67from prisme.workspace import action_table, rms8from tests.conftest import tiny_config9from prisme.model import PrismeModel101112def random_actions(cfg, B, T, R, seed=0, include_stop=True):13 g = torch.Generator().manual_seed(seed)14 hi = cfg.n_actions + (1 if include_stop else 0)15 return torch.randint(0, hi, (B, T, R), generator=g)161718# 1. shapes, dtypes, gradients, finiteness19def test_01_shapes_dtypes_grads_finite(model_nonzero, cfg, tokens):20 m = model_nonzero21 m.train()22 acts = random_actions(cfg, 3, 12, 4)23 out = m(tokens, acts)24 assert out["logits"].shape == (3, 12, cfg.vocab_size)25 assert out["logits"].dtype == torch.float3226 assert torch.isfinite(out["logits"]).all()27 out["logits"].sum().backward()28 for n, p in m.named_parameters():29 if n.startswith("controller"):30 continue31 assert p.grad is not None and torch.isfinite(p.grad).all(), n323334# 2. exact backbone identity with budget 035def test_02_budget_zero_identity(model_nonzero, tokens):36 m = model_nonzero37 ref = m.backbone.head(m.backbone(tokens))38 out = m(tokens, actions=None)["logits"]39 assert torch.equal(out, ref)40 out0 = m(tokens, actions=torch.empty(3, 12, 0, dtype=torch.long))["logits"]41 assert torch.equal(out0, ref)424344# 3. exact identity with immediate STOP (no edits)45def test_03_stop_identity(model_nonzero, cfg, tokens):46 m = model_nonzero47 ref = m.backbone.head(m.backbone(tokens))48 stop = torch.full((3, 12, 4), cfg.stop_action)49 assert torch.equal(m(tokens, stop)["logits"], ref)505152# 4. zero-initialised readout adapter preserves identity even with edits53def test_04_zero_readout_identity(model, cfg, tokens):54 assert torch.count_nonzero(model.readout.U.weight) == 055 ref = model.backbone.head(model.backbone(tokens))56 acts = random_actions(cfg, 3, 12, 8, include_stop=False)57 out = model(tokens, acts, return_state=True)58 assert (out["S_final"] != out["ctx"].S0).any() # edits happened59 assert torch.equal(out["logits"], ref)606162# 5. future suffix changes do not alter preceding logits (with edits, nonzero readout)63def test_05_causality(model_nonzero, cfg, tokens):64 m = model_nonzero65 acts = random_actions(cfg, 3, 12, 4, include_stop=False)66 a = m(tokens, acts)["logits"]67 t2 = tokens.clone()68 t2[:, 8:] = (t2[:, 8:] + 5) % cfg.vocab_size69 b = m(t2, acts)["logits"]70 assert torch.allclose(a[:, :8], b[:, :8], atol=1e-5, rtol=1e-5)71 assert not torch.allclose(a[:, 8:], b[:, 8:])727374# 6. inactive slots exactly unchanged after singles and pairs75def test_06_locality(model_nonzero, cfg, tokens):76 m = model_nonzero77 ctx = m.encode(tokens)78 tbl = action_table(cfg)79 for a in range(cfg.n_actions + 1):80 act = torch.full((3, 12), a)81 S1, info = m.apply_action(ctx.S0, ctx, act, 0, 3)82 slots = [s for s in tbl[a].tolist() if s >= 0]83 for j in range(cfg.n_slots):84 if j in slots:85 assert not torch.equal(S1[:, :, j], ctx.S0[:, :, j]), (a, j)86 else:87 assert torch.equal(S1[:, :, j], ctx.S0[:, :, j]), (a, j)888990# 7. backbone / source-bank states unchanged by workspace editing91def test_07_bank_unchanged(model_nonzero, cfg, tokens):92 m = model_nonzero93 ctx = m.encode(tokens)94 snap = {k: getattr(ctx, k).clone() for k in ("H", "K", "V", "z", "S0")}95 S = ctx.S096 for r in range(4):97 S, _ = m.apply_action(S, ctx, random_actions(cfg, 3, 12, 1, seed=r)[:, :, 0], r, 3 - r)98 for k, v in snap.items():99 assert torch.equal(getattr(ctx, k), v), k100101102# 8. RMS displacement bound: RMS(S_j^n - S_j^0) <= n_j * eta_max * c103def test_08_rms_displacement_bound(model_nonzero, cfg, tokens):104 m = model_nonzero105 ctx = m.encode(tokens)106 S = ctx.S0107 counts = torch.zeros(3, 12, cfg.n_slots)108 for r in range(8):109 act = random_actions(cfg, 3, 12, 1, seed=100 + r)[:, :, 0]110 S, info = m.apply_action(S, ctx, act, r, 7 - r)111 counts += info["write_mask"].float()112 disp = rms(S - ctx.S0) # [B,T,M]113 bound = counts * cfg.eta_max * cfg.clip_c114 assert (disp <= bound + 1e-5).all()115116117# 9. differential correction is zero for identical initial/final states118def test_09_zero_delta_identical_states(model_nonzero, tokens):119 m = model_nonzero120 ctx = m.encode(tokens)121 d = m.readout(ctx.S0, ctx.S0, ctx.z)122 assert torch.equal(d, torch.zeros_like(d))123124125# 10. the hard action limit cannot be exceeded126def test_10_hard_cap(model_nonzero, cfg, tokens):127 with pytest.raises(AssertionError):128 model_nonzero(tokens, random_actions(cfg, 3, 12, cfg.r_max + 1))129 cost = torch.full((cfg.n_actions,), -1.0) # negative cost -> controller always wants to edit130 with torch.no_grad():131 model_nonzero.controller.heads.bias[0] = 5.0132 out = model_nonzero.forward_policy(tokens, max_budget=50, cost=cost)133 assert out["actions"].shape[-1] == cfg.r_max134 assert int(out["n_edits"].max()) <= cfg.r_max135136137# 11. simultaneous pair edits do not depend on lane/iteration order138def test_11_pair_order_independence(model_nonzero, cfg, tokens):139 m = model_nonzero140 ctx = m.encode(tokens)141 pair = torch.full((3, 12), cfg.n_slots) # PAIR_0_1142 S_a, _ = m.apply_action(ctx.S0, ctx, pair, 0, 0)143 tbl = m.editor.actions.clone()144 tbl[cfg.n_slots] = torch.tensor([1, 0]) # swap lanes145 m.editor.actions = tbl146 S_b, _ = m.apply_action(ctx.S0, ctx, pair, 0, 0)147 assert torch.equal(S_a, S_b)148149150# 12. valid gradients reach the readout, and then the editor after bootstrap (U != 0)151def test_12_gradient_flow(cfg, tokens):152 torch.manual_seed(0)153 m = PrismeModel(cfg)154 acts = random_actions(cfg, 3, 12, 2, include_stop=False)155 loss = m(tokens, acts)["logits"].logsumexp(-1).mean()156 loss.backward()157 assert m.readout.U.weight.grad is not None and m.readout.U.weight.grad.abs().sum() > 0158 ed = [p.grad for n, p in m.named_parameters() if n.startswith("editor.block")]159 assert all(g is None or g.abs().sum() == 0 for g in ed) # U = 0 blocks the path, as predicted160 m.zero_grad()161 with torch.no_grad():162 m.readout.U.weight.normal_(std=0.05)163 m(tokens, acts)["logits"].logsumexp(-1).mean().backward()164 ed = [p.grad for n, p in m.named_parameters() if n.startswith("editor.block")]165 assert any(g is not None and g.abs().sum() > 0 for g in ed)166167168# 13. batched teacher forcing == per-sequence forward (masks are per query)169def test_13_batched_equals_single(model_nonzero, cfg, tokens):170 m = model_nonzero171 acts = random_actions(cfg, 3, 12, 3, include_stop=False)172 full = m(tokens, acts)["logits"]173 for b in range(3):174 single = m(tokens[b:b + 1], acts[b:b + 1])["logits"]175 assert torch.allclose(full[b:b + 1], single, atol=1e-5, rtol=1e-5)176177178# 14. no cross-request contamination179def test_14_no_cross_request_contamination(model_nonzero, cfg, tokens):180 m = model_nonzero181 acts = random_actions(cfg, 3, 12, 3, include_stop=False)182 a = m(tokens, acts, return_state=True)183 t2 = tokens.clone()184 t2[1] = (t2[1] + 3) % cfg.vocab_size185 b = m(t2, acts, return_state=True)186 assert torch.allclose(a["logits"][0], b["logits"][0], atol=1e-5)187 assert torch.allclose(a["logits"][2], b["logits"][2], atol=1e-5)188 assert torch.equal(a["S_final"][0], b["S_final"][0])189190191# 15. determinism under fixed seed192def test_15_determinism(cfg, tokens):193 outs = []194 for _ in range(2):195 torch.manual_seed(99)196 m = PrismeModel(cfg)197 with torch.no_grad():198 m.readout.U.weight.normal_(std=0.05)199 outs.append(m(tokens, random_actions(cfg, 3, 12, 4))["logits"])200 assert torch.equal(outs[0], outs[1])201202203# controller features are target-free and never execute edits204def test_controller_target_free(model_nonzero, cfg, tokens):205 m = model_nonzero206 ctx = m.encode(tokens)207 step = torch.zeros(3, 12, dtype=torch.long)208 p1 = m.controller(ctx.z, ctx.S0, step, step + 3)209 t2 = tokens.clone(); t2[:, -1] = (t2[:, -1] + 1) % cfg.vocab_size # change only the last token210 ctx2 = m.encode(t2)211 p2 = m.controller(ctx2.z, ctx2.S0, step, step + 3)212 assert torch.allclose(p1["g_hat"][:, :-1], p2["g_hat"][:, :-1], atol=1e-5)213 assert p1["g_hat"].shape == (3, 12, cfg.n_actions)214215216def test_dense_writes_ablation_touches_all_slots(tokens):217 cfg = tiny_config(dense_writes=True)218 torch.manual_seed(0)219 m = PrismeModel(cfg)220 ctx = m.encode(tokens)221 S1, info = m.apply_action(ctx.S0, ctx, torch.zeros(3, 12, dtype=torch.long), 0, 0)222 assert (S1 != ctx.S0).any(-1).all()223 S2, _ = m.apply_action(ctx.S0, ctx, torch.full((3, 12), cfg.stop_action), 0, 0)224 assert torch.equal(S2, ctx.S0)225226227def test_controller_h_entropy_features_target_free(tokens):228 cfg = tiny_config(ctrl_input="h", ctrl_use_entropy=True)229 torch.manual_seed(0)230 m = PrismeModel(cfg).eval()231 ctx = m.encode(tokens)232 assert ctx.ctrl_in.shape[-1] == cfg.d_model and ctx.entropy.shape == tokens.shape233 step = torch.zeros(3, 12, dtype=torch.long)234 p1 = m.controller(ctx.ctrl_in, ctx.S0, step, step + 3, ctx.entropy)235 t2 = tokens.clone(); t2[:, -1] = (t2[:, -1] + 1) % cfg.vocab_size236 ctx2 = m.encode(t2)237 p2 = m.controller(ctx2.ctrl_in, ctx2.S0, step, step + 3, ctx2.entropy)238 assert torch.allclose(p1["g_hat"][:, :-1], p2["g_hat"][:, :-1], atol=1e-5)239 cost = torch.full((cfg.n_actions,), 0.05)240 out = m.forward_policy(tokens, 4, cost)241 assert out["logits"].shape == (3, 12, cfg.vocab_size)242