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
|
import torch
import torch.nn.functional as F
from src.config import Config
from src.data import build_examples, make_batch
from src.model_api import build_model
def _batch(cfg, inputs=(2, 7, 42, 99)):
ex = build_examples(list(inputs), cfg)
return make_batch(ex, cfg)
def test_rnn_shapes_and_halting_range():
cfg = Config(model="rnn", max_steps=20, min_steps=2)
m = build_model(cfg)
b = _batch(cfg)
out = m(b["x"], b["y_in"])
assert out["logits"].shape == (4, b["y"].shape[1], cfg.vocab)
s = out["halt_steps"]
assert s.shape == (4,)
assert bool((s >= cfg.min_steps).all())
assert bool((s <= cfg.max_steps).all())
def test_rnn_fixed_k_ablation():
cfg = Config(model="rnn", halting=False, max_steps=20)
m = build_model(cfg)
b = _batch(cfg)
out = m(b["x"], b["y_in"])
assert torch.allclose(out["halt_steps"], torch.full((4,), 20.0))
def test_transformer_shapes():
cfg = Config(model="transformer")
m = build_model(cfg)
b = _batch(cfg)
out = m(b["x"], b["y_in"])
assert out["logits"].shape == (4, b["y"].shape[1], cfg.vocab)
assert out["halt_steps"] is None
def test_transformer_integers_mode():
cfg = Config(model="transformer", vocab_mode="integers")
m = build_model(cfg)
b = _batch(cfg)
out = m(b["x"], b["y_in"])
assert out["logits"].shape == (4, b["y"].shape[1], cfg.vocab)
def test_rnn_fully_tied_no_gru_decoder():
"""Regression: design review fix — everything recurrent must be the ONE tied cell."""
import torch.nn as nn
cfg = Config(model="rnn")
m = build_model(cfg)
assert not any(isinstance(mod, nn.GRUCell) for mod in m.modules())
assert not any(isinstance(mod, (nn.GRU, nn.LSTM, nn.RNN)) for mod in m.modules())
def test_param_counts_logged_not_gated():
r = build_model(Config(model="rnn")).param_count()
t = build_model(Config(model="transformer")).param_count()
assert r > 1000 and t > 1000
def test_forward_backward_no_nan_both_models():
for model_name in ("rnn", "transformer"):
cfg = Config(model=model_name)
m = build_model(cfg)
b = _batch(cfg)
out = m(b["x"], b["y_in"])
y_safe = b["y"].clamp(max=cfg.vocab - 1)
loss = F.cross_entropy(out["logits"].reshape(-1, cfg.vocab), y_safe.reshape(-1), reduction="none")
mask = b["y_mask"].float().reshape(-1)
loss = (loss * mask).sum() / mask.sum()
loss.backward()
assert torch.isfinite(loss).item(), f"{model_name} loss NaN"
def test_greedy_decode_shape():
from src.model_api import greedy_decode
cfg = Config(model="rnn")
m = build_model(cfg)
x = torch.tensor([[4, 2], [9, 9]], dtype=torch.long)
gen = greedy_decode(m, x, cfg)
assert gen.shape == (2, cfg.max_out_len)
|
