1 files changed, 43 insertions, 0 deletions
diff --git a/tests/test_batched_eval.py b/tests/test_batched_eval.py
new file mode 100644
index 0000000..2f09c82
--- /dev/null
+++ b/tests/test_batched_eval.py
@@ -0,0 +1,43 @@
+"""Batched evaluate() correctness — the speedup rewrite must preserve exact semantics."""
+import torch
+
+from src.config import Config
+from src.data import build_examples, get_splits
+from src.train import evaluate
+
+from tests.test_eval_classification import StubModel, _sieve35
+
+
+def test_batched_evaluate_known_model():
+    """The {2,3,5,7} sieve is CORRECT on the whole [2,100] range (any composite <= 101 has a
+    factor in {2,3,5,7}, and every in-range next prime has no small divisor) — so evaluate()
+    against the sieve stub must return token acc 1.0 and exact-match 1.0 on both splits."""
+    cfg = Config()
+    train_in, val_in = get_splits(cfg)
+    m = StubModel(_sieve35)
+    for split_in in (train_in, val_in):
+        ex = build_examples(split_in, cfg)
+        tok, em, per = evaluate(m, ex, cfg)
+        assert tok == 1.0, f"token acc {tok} (split size {len(ex)})"
+        assert em == 1.0, f"exact-match {em} (split size {len(ex)})"
+        assert len(per) == len(ex)
+        assert all(ok for _, _, _, ok in per)
+
+
+def test_batched_evaluate_matches_singleton_greedy():
+    """For a REAL model, batched evaluate's greedy decode must equal greedy_decode per row
+    (batch-invariance guarantees this — belt and braces against a regressed layout fix)."""
+    from src.model_api import build_model, greedy_decode
+    from src.data import make_batch, decode_tokens
+    cfg = Config(model="rnn")
+    torch.manual_seed(0)
+    m = build_model(cfg)
+    _, val_in = get_splits(cfg)
+    val_ex = build_examples(val_in, cfg)
+    _, em, per = evaluate(m, val_ex, cfg)
+    # independent per-example greedy via the public helper
+    for i, (x, target) in enumerate(val_ex):
+        xi = torch.tensor(x, dtype=torch.long).unsqueeze(0)
+        gen = greedy_decode(m, xi, cfg)[0].tolist()
+        pred_alone = decode_tokens(gen, cfg)
+        assert per[i][2] == pred_alone, f"example {i}: batched {per[i][2]} vs singleton {pred_alone}"