1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
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())
|
