verify.py

 1"""Numerical, shape, gradient, masking, and causality checks for the course."""
 2import numpy as np
 3import torch
 4import numpy_core as n
 5import torch_core as t
 6from train_copy import dataset
 7
 8
 9def check() -> None:
10    torch.set_num_threads(1)
11    rng = np.random.default_rng(19)
12    q, k, v = rng.normal(size=(2, 3, 4, 8)), rng.normal(size=(2, 3, 6, 8)), rng.normal(size=(2, 3, 6, 5))
13    visible = np.ones((2, 1, 4, 6), dtype=bool)
14    visible[..., -1] = False
15    result, weights = n.attention(q, k, v, visible)
16    tensors = [torch.tensor(x, requires_grad=True) for x in (q, k, v)]
17    torch_result, torch_weights = t.attention(tensors[0], tensors[1], tensors[2], torch.tensor(visible))
18    np.testing.assert_allclose(result, torch_result.detach().numpy(), atol=1e-12)
19    np.testing.assert_allclose(weights.sum(-1), 1.0)
20    assert np.all(weights[..., -1] == 0)
21    changed_v = v.copy(); changed_v[..., -1, :] += 10000
22    np.testing.assert_allclose(n.attention(q, k, changed_v, visible)[0], result)
23    torch_result.square().sum().backward()
24    assert all(x.grad is not None and torch.isfinite(x.grad).all() for x in tensors)
25    for implementation, data, mask in [(n, (q,k,v), np.zeros_like(visible)),
26                                     (t, tensors, torch.zeros_like(torch.tensor(visible)))]:
27        try:
28            implementation.attention(*data, allowed=mask)
29        except ValueError:
30            pass
31        else:
32            raise AssertionError("All-masked rows must be rejected")
33
34    np.testing.assert_allclose(n.positions(9,8), t.positions(9,8,dtype=torch.float64).numpy(), atol=1e-12)
35    np.testing.assert_allclose(n.embedding_demo()[2], t.embedding_demo()[2].detach().numpy())
36    x = rng.normal(size=(2,4,8))
37    np.testing.assert_allclose(n.layer_norm(x), t.layer_norm(torch.tensor(x)).numpy(), atol=1e-12)
38    mha_n, mha_t = n.MultiHead(8,2,seed=3), t.MultiHead(8,2,seed=3).double()
39    # Copy full precision arrays; converting previously rounded float32 weights to double cannot undo rounding.
40    with torch.no_grad():
41        for name in ['wq','wk','wv','wo']:
42            getattr(mha_t,name).weight.copy_(torch.tensor(getattr(mha_n,name).T))
43    np.testing.assert_allclose(mha_n(x,x)[0], mha_t(torch.tensor(x),torch.tensor(x))[0].detach().numpy(), atol=1e-12)
44    permutation = np.array([2,0,3,1])
45    np.testing.assert_allclose(mha_n(x[:,permutation],x[:,permutation])[0],mha_n(x,x)[0][:,permutation],atol=1e-12)
46
47    source = torch.tensor([[4,7,0],[9,4,6]])
48    prefix = torch.tensor([[1,4,7,0],[1,9,4,6]])
49    model = t.TinyTransformer().eval()
50    with torch.no_grad():
51        baseline = model(source,prefix)
52        changed = prefix.clone(); changed[:,3] = 11
53        torch.testing.assert_close(model(source,changed)[:,:3],baseline[:,:3])
54        padded_source = torch.cat([source,torch.zeros(2,2,dtype=torch.long)],dim=1)
55        torch.testing.assert_close(model(padded_source,prefix),baseline,atol=2e-6,rtol=2e-6)
56        memory, mask = model.encode(source)
57        for i in range(prefix.shape[1]):
58            incremental = model.decode(prefix[:,:i+1],memory,mask)[:,-1]
59            torch.testing.assert_close(incremental,baseline[:,i],atol=2e-6,rtol=2e-6)
60    train, held_out = dataset()
61    assert not set(map(tuple,train)) & set(map(tuple,held_out))
62    assert n.demo_forward().shape == (2,3,12)
63    print("PASS: NumPy/PyTorch agreement, normalized weights, masked-value invariance, gradients,")
64    print("positions, embeddings, normalization, heads, permutation equivariance, decoder causality,")
65    print("padding invariance, incremental-prefix equivalence, and disjoint evaluation sequences.")
66
67
68if __name__ == "__main__":
69    check()