tests/test_transformer.py

1import numpy as np23from optim import Adam4from transformer import (ModelConfig, init_params, count_params, forward, loss_and_grads,5                         embed_forward, embed_backward, ffn_forward, split_heads,6                         merge_heads, attention_forward)789def tiny_cfg(**kw):10    base = dict(vocab_size=7, context_len=8, emb_dim=4, num_heads=2, ff_dim=8, num_layers=2)11    base.update(kw)12    return ModelConfig(**base)131415def test_param_count_matches_formula():16    # #params = (2d+1)V + (Tmax+2)d + L(4d^2 + 2dF + F + 9d); d=32, F=128, L=2, Tmax=32, V=30017    cfg = ModelConfig(300, 32, 32, 2, 128, 2)18    params = init_params(cfg, np.random.default_rng(0))19    assert count_params(params) == 45_996202122def test_initial_loss_is_close_to_ln_vocab():23    cfg = ModelConfig(300, 32, 32, 2, 128, 2)24    rng = np.random.default_rng(0)25    params = init_params(cfg, rng)26    ids, targets = rng.integers(300, size=16), rng.integers(300, size=16)27    loss, _ = loss_and_grads(params, ids, targets, cfg)28    assert abs(loss - np.log(300)) < 0.5293031def test_embedding_worked_example():32    params = {"tok_emb": np.array([[.1, .2], [.3, -.1], [.5, .4], [-.2, .6]]),33              "pos_emb": np.array([[0, .1], [.1, 0], [.2, .1]])}34    X0 = embed_forward(params, np.array([2, 0, 2]))35    assert np.allclose(X0, [[.5, .5], [.2, .2], [.7, .5]])363738def test_embedding_backward_sums_repeated_ids():39    grads = {"tok_emb": np.zeros((4, 2)), "pos_emb": np.zeros((3, 2))}40    dX0 = np.array([[1., 0.], [0., 1.], [2., 2.]])41    embed_backward(dX0, np.array([2, 0, 2]), grads)42    assert np.allclose(grads["tok_emb"][2], [3, 2])       # rows 0 and 2 both landed here43    assert np.allclose(grads["tok_emb"][0], [0, 1])44    assert np.allclose(grads["pos_emb"], dX0)454647def test_ffn_worked_example():48    p = "b0."49    params = {p + "ff1_w": np.array([[1., 0., -1.], [0., 1., 1.]]),50              p + "ff1_b": np.array([0., .5, 1.]),51              p + "ff2_w": np.array([[1., 2.], [0., 1.], [1., -1.]]),52              p + "ff2_b": np.zeros(2)}53    y, cache = ffn_forward(np.array([[1., -1.]]), params, p)54    assert np.allclose(cache["Z"], [[1, -.5, -1]])55    assert np.allclose(cache["A"], [[1, 0, 0]])           # ReLU zeroes the negatives56    assert np.allclose(y, [[1, 2]])575859def test_split_merge_heads_round_trip():60    M = np.arange(24, dtype=float).reshape(3, 8)61    Mh = split_heads(M, 2)62    assert Mh.shape == (2, 3, 4)63    assert np.array_equal(merge_heads(Mh), M)646566def test_attention_weights_are_causal_and_normalised():67    cfg = tiny_cfg(num_layers=1)68    params = init_params(cfg, np.random.default_rng(0))69    Xn = np.random.default_rng(1).normal(size=(5, 4))70    _, c = attention_forward(Xn, params, "b0.", cfg)71    A = c["A"]                                            # (H, T, T)72    assert np.allclose(np.triu(A[0], 1), 0.0)             # nothing flows from the future73    assert np.allclose(A.sum(axis=-1), 1.0)74    assert np.allclose(A[:, 0, 0], 1.0)                   # first token only sees itself757677def test_future_tokens_do_not_change_earlier_logits():78    cfg = tiny_cfg()79    params = init_params(cfg, np.random.default_rng(0))80    a = np.array([1, 2, 3, 4, 5])81    b = a.copy()82    b[4] = 0                                              # change only the last token83    la, _ = forward(params, a, cfg)84    lb, _ = forward(params, b, cfg)85    assert np.allclose(la[:4], lb[:4])86    assert not np.allclose(la[4], lb[4])878889def test_grads_have_the_shape_of_params():90    cfg = tiny_cfg()91    rng = np.random.default_rng(0)92    params = init_params(cfg, rng)93    _, grads = loss_and_grads(params, rng.integers(7, size=5), rng.integers(7, size=5), cfg)94    assert grads.keys() == params.keys()95    assert all(grads[k].shape == params[k].shape for k in params)969798def test_overfit_one_window():99    # a correct forward + backward + optimiser must be able to memorise 8 tokens100    cfg = tiny_cfg(emb_dim=16, ff_dim=32)101    rng = np.random.default_rng(0)102    params = init_params(cfg, rng)103    ids, targets = rng.integers(7, size=8), rng.integers(7, size=8)104    opt = Adam(params)105    for _ in range(200):106        loss, grads = loss_and_grads(params, ids, targets, cfg)107        opt.step(params, grads, lr=1e-2)108    # (no "starts near ln V" check here: with V = 7 the initial logit variance109    #  ~ 2d/(d+V) is too large for that rule; the V = 300 test covers it)110    assert loss < 0.05