implement Experiment 1: data pipeline, tied-RNN w/ ACT halting, transformer baseline, train/eval/plot, 16 tests

2 files changed, 155 insertions, 0 deletions

diff --git a/tests/test_data.py b/tests/test_data.py
new file mode 100644
index 0000000..fe84e94
--- /dev/null
+++ b/tests/test_data.py
@@ -0,0 +1,78 @@
+import torch
+
+from src.config import Config
+from src.data import (build_examples, decode_tokens, encode_int, get_splits,
+                      make_batch, next_prime, sieve_primes)
+
+
+def test_sieve_primes_known():
+    assert sieve_primes(30) == [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
+    assert sieve_primes(1) == []
+    assert sieve_primes(2) == [2]
+
+
+def test_next_prime():
+    p = sieve_primes(200)
+    assert next_prime(42, p) == 43
+    assert next_prime(100, p) == 101
+    assert next_prime(2, p) == 3
+    assert next_prime(97, p) == 101
+
+
+def test_encode_decode_roundtrip_digits():
+    cfg = Config(vocab_mode="digits")
+    for n in [2, 7, 10, 42, 99, 100]:
+        assert decode_tokens(encode_int(n, cfg) + [cfg.eos_id], cfg) == n
+    assert encode_int(42, cfg) == [4, 2]
+
+
+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
+
+
+def test_splits_sizes_and_no_overlap():
+    cfg = Config()
+    tr, va = get_splits(cfg)
+    assert len(tr) == 69 and len(va) == 30
+    assert not set(tr) & set(va)
+    assert set(tr) | set(va) == set(range(2, 101))
+
+
+def test_splits_seed_stable():
+    a1, b1 = get_splits(Config(seed=3))
+    a2, b2 = get_splits(Config(seed=3))
+    assert a1 == a2 and b1 == b2
+    a3, _ = get_splits(Config(seed=4))
+    assert a1 != a3
+
+
+def test_build_examples_targets():
+    cfg = Config()
+    ex = build_examples([2, 42, 99, 100], cfg)
+    assert [decode_tokens(y, cfg) for _, y in ex] == [3, 43, 101, 101]
+    assert ex[0][1][-1] == cfg.eos_id  # EOS-terminated
+
+
+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["y_in"].shape == (4, 4)
+    assert b["y_mask"].shape == (4, 4)
+    # pad positions masked out
+    assert not b["y_mask"][0, 3]           # "3"+EOS row has pad at index 3
+    assert b["y_mask"][0, 1]
+    # BOS position of y_in is eos_id
+    assert (b["y_in"][:, 0] == cfg.eos_id).all()
+
+
+def test_flagged_composites_are_composite():
+    # sanity: the diagnostic set in eval.py really is composite and needs divisors > 7
+    for n in (121, 143, 169, 187):
+        assert any(n % d == 0 for d in range(2, int(n ** 0.5) + 1))
+        assert all(n % d != 0 for d in (2, 3, 5, 7))
diff --git a/tests/test_models.py b/tests/test_models.py
new file mode 100644
index 0000000..7994e86
--- /dev/null
+++ b/tests/test_models.py
@@ -0,0 +1,77 @@
+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_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)