fix codex BLOCKERs: global fixed layout invariance, min_steps off-by-one, integers EOS alias, rerun guard, eval prereg-literal codes + dual-checkpoint honesty; +7 regression tests

2 files changed, 110 insertions, 3 deletions

diff --git a/tests/test_codex_fixes.py b/tests/test_codex_fixes.py
new file mode 100644
index 0000000..9e14113
--- /dev/null
+++ b/tests/test_codex_fixes.py
@@ -0,0 +1,104 @@
+"""Regression tests for codex-review fixes (design/reviews/codex-review.md)."""
+import torch
+
+from src.config import Config
+from src.data import build_examples, decode_tokens, make_batch
+from src.model_api import build_model, greedy_decode
+
+
+def test_mixed_batch_vs_singleton_invariance_rnn():
+    cfg = Config(model="rnn")
+    m = build_model(cfg)
+    ex = build_examples([2, 42, 99], cfg)
+    batch = make_batch(ex, cfg)
+    out_batch = m(batch["x"], batch["y_in"])
+    for i in range(3):
+        single = make_batch([ex[i]], cfg)
+        out_single = m(single["x"], single["y_in"])
+        assert torch.allclose(out_batch["logits"][i], out_single["logits"][0], atol=1e-6), f"row {i}"
+        assert torch.allclose(out_batch["halt_steps"][i], out_single["halt_steps"][0], atol=1e-6)
+
+
+def test_mixed_batch_vs_singleton_invariance_transformer():
+    cfg = Config(model="transformer")
+    m = build_model(cfg)
+    ex = build_examples([2, 42, 99], cfg)
+    batch = make_batch(ex, cfg)
+    out_batch = m(batch["x"], batch["y_in"])
+    for i in range(3):
+        single = make_batch([ex[i]], cfg)
+        out_single = m(single["x"], single["y_in"])
+        assert torch.allclose(out_batch["logits"][i], out_single["logits"][0], atol=1e-6), f"row {i}"
+
+
+def test_greedy_decode_layout_invariant():
+    cfg = Config(model="transformer")
+    m = build_model(cfg)
+    ex = build_examples([42], cfg)
+    single = make_batch(ex, cfg)
+    out_a = greedy_decode(m, single["x"], cfg)
+    raw = torch.tensor([[4, 2]], dtype=torch.long)   # unpadded — greedy must left-pad internally
+    out_b = greedy_decode(m, raw, cfg)
+    assert torch.equal(out_a, out_b)
+
+
+def test_min_steps_floor_reachable():
+    """min_steps off-by-one fix: earliest halt is step min_steps (forced-1 halt prob -> exactly floor)."""
+    cfg = Config(model="rnn", min_steps=2)
+    m = build_model(cfg)
+    with torch.no_grad():
+        m.halt_head.bias.fill_(50.0)   # p -> 1 at the first free step
+    ex = build_examples([2, 7, 42], cfg)
+    b = make_batch(ex, cfg)
+    out = m(b["x"], b["y_in"])
+    assert torch.allclose(out["halt_steps"], torch.full_like(out["halt_steps"], float(cfg.min_steps)), atol=1e-3)
+
+
+def test_pad_logits_cannot_change_loss():
+    """Loss and grads must be insensitive to logits at padded target positions."""
+    cfg = Config(model="rnn")
+    m = build_model(cfg)
+    b = make_batch(build_examples([2, 7], cfg), cfg)
+    out = m(b["x"], b["y_in"])
+    y_safe = b["y"].clamp(max=cfg.vocab - 1)
+    loss = torch.nn.functional.cross_entropy(out["logits"].reshape(-1, cfg.vocab),
+                                             y_safe.reshape(-1), reduction="none")
+    mask = b["y_mask"].float().reshape(-1)
+    l1 = (loss * mask).sum() / mask.sum()
+    l1.backward()
+    grads = {k: v.grad.clone() for k, v in m.named_parameters() if v.grad is not None}
+    m.zero_grad()
+    out2 = m(b["x"], b["y_in"])
+    logits2 = out2["logits"].clone()
+    pad_positions = ~b["y_mask"]                       # (B,T_out)
+    logits2[pad_positions] += 100.0                    # huge perturbation on pad positions only
+    loss2 = torch.nn.functional.cross_entropy(logits2.reshape(-1, cfg.vocab),
+                                              y_safe.reshape(-1), reduction="none")
+    l2 = (loss2 * mask).sum() / mask.sum()
+    assert torch.allclose(l1.detach(), l2.detach())
+    l2.backward()
+    for k, v in m.named_parameters():
+        if v.grad is not None:
+            assert torch.allclose(grads[k], v.grad, atol=1e-6), k
+
+
+def test_integers_vocab_eos_not_aliased():
+    """Value 101 (= next_prime(100)) must not collide with EOS in integers mode."""
+    cfg = Config(vocab_mode="integers", range_end=100)
+    ex = build_examples([100], cfg)
+    _, target = ex[0]
+    assert target[0] == 101 and target[1] == cfg.eos_id
+    assert cfg.eos_id != 101
+    assert cfg.vocab == 103
+    assert decode_tokens(target, cfg) == 101
+
+
+def test_act_weights_sum_to_one():
+    """ACT mass conservation: steps must always lie in [min_steps, K] for random states."""
+    cfg = Config(model="rnn", max_steps=20, min_steps=2)
+    m = build_model(cfg)
+    b = make_batch(build_examples([2, 7, 42, 99, 3, 15, 60, 97], cfg), cfg)
+    out = m(b["x"], b["y_in"])
+    s = out["halt_steps"]
+    assert bool((s >= cfg.min_steps).all())
+    assert bool((s <= cfg.max_steps).all())
diff --git a/tests/test_data.py b/tests/test_data.py
index fe84e94..cd4e620 100644
--- a/tests/test_data.py
+++ b/tests/test_data.py
@@ -30,7 +30,7 @@ def test_encode_decode_roundtrip_integers():
     cfg = Config(vocab_mode="integers", range_end=100)
     for n in [2, 42, 100]:
         assert decode_tokens(encode_int(n, cfg) + [cfg.eos_id], cfg) == n
-    assert cfg.vocab == 102  # 0..101 + EOS
+    assert cfg.vocab == 103  # values 0..101, EOS=102, pad=103 (no alias with target 101)
 
 
 def test_splits_sizes_and_no_overlap():
@@ -60,8 +60,8 @@ def test_make_batch_shapes_and_padding():
     cfg = Config()
     ex = build_examples([2, 7, 42, 99], cfg)
     b = make_batch(ex, cfg)
-    assert b["x"].shape == (4, 2)          # max input len 2 (42, 99)
-    assert b["y"].shape == (4, 4)          # max output len 4 ("101"+EOS)
+    assert b["x"].shape == (4, 3)          # GLOBAL layout: left-padded to 3 digits (max of [2,100])
+    assert b["y"].shape == (4, 4)          # GLOBAL layout: right-padded to 4 ("101"+EOS)
     assert b["y_in"].shape == (4, 4)
     assert b["y_mask"].shape == (4, 4)
     # pad positions masked out
@@ -69,6 +69,9 @@ def test_make_batch_shapes_and_padding():
     assert b["y_mask"][0, 1]
     # BOS position of y_in is eos_id
     assert (b["y_in"][:, 0] == cfg.eos_id).all()
+    # LEFT-padding: input 2 sits at the last column, pads at the front
+    assert b["x"][0, 2] == 2 and b["x"][0, 0] == cfg.pad_id
+    assert b["x"][2, 1] == 4 and b["x"][2, 2] == 2
 
 
 def test_flagged_composites_are_composite():