1 files changed, 148 insertions, 0 deletions
diff --git a/tests/test_scratchpad.py b/tests/test_scratchpad.py
new file mode 100644
index 0000000..2103453
--- /dev/null
+++ b/tests/test_scratchpad.py
@@ -0,0 +1,148 @@
+"""Unit tests for Addendum 8: 6-arm token-space recurrence & decomposition."""
+import pytest
+import torch
+import torch.nn as nn
+
+from src.config import Config
+from src.data import (
+    EOS_ID,
+    NOISE_SLOT_ID,
+    PAUSE_ID,
+    RAND_START_ID,
+    SEP_ID,
+    SYM_C,
+    SYM_COLON,
+    SYM_D,
+    SYM_EQ,
+    build_examples,
+    decode_tokens,
+    make_batch,
+    make_structured_trace,
+    next_prime,
+    sieve_primes,
+)
+from src.models.transformer import TransformerBaseline
+
+
+def test_vocab_and_special_tokens():
+    cfg_none = Config(scratch_mode="none")
+    assert cfg_none.vocab == 11
+    assert cfg_none.eos_id == 10
+    assert cfg_none.pad_id == 11
+
+    cfg_scratch = Config(scratch_mode="filler")
+    assert cfg_scratch.vocab == 34
+    assert cfg_scratch.eos_id == 10
+    assert cfg_scratch.pad_id == 34
+
+
+def test_make_structured_trace():
+    primes = sieve_primes(50)
+    # n=14 -> next prime 17. Candidates 15 (comp: d3:0), 16 (comp: d2:0), 17 (prime: d2:1, d3:1)
+    trace = make_structured_trace(14, 17, primes)
+    assert SYM_C in trace
+    assert SYM_EQ in trace
+    assert SYM_D in trace
+    assert SYM_COLON in trace
+
+    # n=2 -> next prime 3. Candidate 3 has no primes with p*p <= 3, trace should be [SYM_C, SYM_EQ, 3]
+    trace_2_3 = make_structured_trace(2, 3, primes)
+    assert trace_2_3 == [SYM_C, SYM_EQ, 3]
+
+
+def test_decode_tokens_all_modes():
+    cfg_none = Config(scratch_mode="none")
+    assert decode_tokens([4, 3, EOS_ID], cfg_none) == 43
+
+    cfg_filler = Config(scratch_mode="filler")
+    # with separator #
+    assert decode_tokens([PAUSE_ID, PAUSE_ID, SEP_ID, 4, 3, EOS_ID], cfg_filler) == 43
+    # without separator (should fail / return -1)
+    assert decode_tokens([4, 3, EOS_ID], cfg_filler) == -1
+
+
+def test_build_examples_and_loss_mask():
+    modes = ["none", "structured", "filler", "random_learned", "random_frozen", "random_noise"]
+    for mode in modes:
+        cfg = Config(scratch_mode=mode, range_start=2, range_end=30, scratch_len=16)
+        exs = build_examples([14], cfg)
+        assert len(exs) == 1
+        x, y = exs[0]
+        batch = make_batch(exs, cfg)
+        loss_mask = batch["loss_mask"][0]
+
+        if mode == "none":
+            assert decode_tokens(y, cfg) == 17
+            assert all(loss_mask[: len(y)])
+        elif mode == "structured":
+            assert SEP_ID in y
+            assert decode_tokens(y, cfg) == 17
+            assert all(loss_mask[: len(y)])  # structured trace is fully supervised
+        elif mode in ("filler", "random_learned", "random_frozen", "random_noise"):
+            assert SEP_ID in y
+            assert decode_tokens(y, cfg) == 17
+            # first 16 tokens must have loss_mask == 0
+            assert all(not m for m in loss_mask[:16])
+            # tokens after separator must have loss_mask == 1
+            assert all(m for m in loss_mask[16 : len(y)])
+
+
+def test_random_tokens_no_digit_leakage():
+    for mode in ("random_learned", "random_frozen"):
+        cfg = Config(scratch_mode=mode, range_start=2, range_end=100, scratch_len=16)
+        exs = build_examples(list(range(2, 101)), cfg)
+        for _, y in exs:
+            random_segment = y[:16]
+            for tok in random_segment:
+                assert RAND_START_ID <= tok <= RAND_START_ID + 15
+                assert not (0 <= tok <= 9)  # no digits leakage
+
+
+def test_transformer_forward_all_arms():
+    modes = ["none", "structured", "filler", "random_learned", "random_frozen", "random_noise"]
+    for mode in modes:
+        cfg = Config(scratch_mode=mode, range_start=2, range_end=30, d_model=32, n_layers=1, n_heads=2)
+        exs = build_examples([14, 15, 16], cfg)
+        batch = make_batch(exs, cfg)
+        model = TransformerBaseline(cfg)
+        out = model(batch["x"], batch["y_in"])
+        logits = out["logits"]
+        assert logits.shape == (3, batch["y"].shape[1], cfg.vocab)
+
+        # compute masked loss
+        ce = nn.CrossEntropyLoss(reduction="none")
+        y_safe = batch["y"].clamp(max=cfg.vocab - 1)
+        loss_tokens = ce(logits.reshape(-1, cfg.vocab), y_safe.reshape(-1)).reshape(3, -1) * batch["loss_mask"].float()
+        loss = loss_tokens.sum() / batch["loss_mask"].sum().clamp(min=1)
+        assert not torch.isnan(loss)
+        loss.backward()
+
+
+def test_random_frozen_zero_gradients():
+    cfg = Config(scratch_mode="random_frozen", range_start=2, range_end=30, d_model=32, n_layers=1, n_heads=2)
+    exs = build_examples([14, 15], cfg)
+    batch = make_batch(exs, cfg)
+    model = TransformerBaseline(cfg)
+    out = model(batch["x"], batch["y_in"])
+    logits = out["logits"]
+    loss = logits.sum()
+    loss.backward()
+
+    # random token embeddings (13..28) must have zero grad because they were detached
+    grad = model.embed.weight.grad
+    assert grad[RAND_START_ID : RAND_START_ID + 16].abs().sum() == 0
+    # non-random tokens must have non-zero grad
+    assert grad[:10].abs().sum() > 0
+
+
+def test_random_noise_stochasticity():
+    cfg = Config(scratch_mode="random_noise", range_start=2, range_end=30, d_model=32, n_layers=1, n_heads=2)
+    exs = build_examples([14], cfg)
+    batch = make_batch(exs, cfg)
+    model = TransformerBaseline(cfg)
+    model.eval()
+
+    out1 = model(batch["x"], batch["y_in"])["logits"]
+    out2 = model(batch["x"], batch["y_in"])["logits"]
+    # Dynamic Gaussian noise should produce distinct logits between forward passes
+    assert not torch.allclose(out1, out2)