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
|
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 == 103 # values 0..101, EOS=102, pad=103 (no alias with target 101)
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, 3) # GLOBAL layout: left-padded to 3 digits (max of [2,100])
assert b["y"].shape == (4, 4) # GLOBAL layout: right-padded to 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()
# LEFT-padding: input 2 sits at the last column, pads at the front
assert b["x"][0, 2] == 2 and b["x"][0, 0] == cfg.pad_id
assert b["x"][2, 1] == 4 and b["x"][2, 2] == 2
def test_flagged_composites_are_composite():
# sanity: the rank-4 sieve signature set is composite, has no divisors <= 7
from src.eval import _sieve_rank_signature
sig = _sieve_rank_signature(4, 101, 200)
assert sig == {121, 143, 169, 187}, f"unexpected rank-4 signature: {sig}"
for n in sig:
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))
# 209 = 11*19 is also rank-4 but outside [101,200]; verify it appears in wider range
sig_wide = _sieve_rank_signature(4, 101, 210)
assert 209 in sig_wide
|
