add a few tests

README.md

@@ -186,7 +186,12 @@ I've published a number of guides that might contain helpful information, most r
 │   ├── mmlu.py                     # Multiple choice questions, broad topics
 │   └── smoltalk.py                 # Conglomerate dataset of SmolTalk from HF
 ├── tests
-│   └── test_engine.py
+│   ├── test_attention_fallback.py  # FA3/SDPA attention fallback
+│   ├── test_engine.py              # Inference engine, KV cache
+│   ├── test_execution.py           # Sandboxed code execution
+│   ├── test_optim.py               # MuonAdamW optimizer (needs GPU)
+│   ├── test_tasks.py               # Task slicing, mixtures, HubDataset
+│   └── test_tokenizer.py           # BPE round-trips, chat rendering
 └── uv.lock
 ```
 

tests/test_execution.py

@@ -0,0 +1,83 @@
+"""
+Test the sandboxed Python execution used by tool use / HumanEval.
+Each test adversarially checks one property of the sandbox.
+
+python -m pytest tests/test_execution.py -v
+"""
+
+import os
+import time
+import pytest
+from nanochat.execution import execute_code
+
+
+def test_happy_path():
+    result = execute_code("print('hello world')")
+    assert result.success
+    assert result.stdout == "hello world\n"
+
+
+def test_exception_is_captured():
+    result = execute_code("1/0")
+    assert not result.success
+    assert "ZeroDivisionError" in result.error
+
+
+def test_timeout_kills_infinite_loop():
+    start = time.time()
+    result = execute_code("while True: pass", timeout=2.0)
+    elapsed = time.time() - start
+    assert result.timeout and not result.success
+    assert elapsed < 5, f"process was not killed promptly ({elapsed:.1f}s)"
+
+
+def test_memory_limit():
+    # 1GB allocation against the 256MB default limit
+    result = execute_code("x = bytearray(1024 * 1024 * 1024)")
+    assert not result.success
+    assert result.memory_exceeded
+
+
+@pytest.mark.parametrize("evil", [
+    "import os; os.system('echo pwned')",
+    "import shutil; shutil.rmtree('/tmp')",
+    "import subprocess; subprocess.Popen(['ls'])",
+    "import os; os.kill(1, 9)",
+    "import os; os.fork()",
+])
+def test_destructive_functions_disabled(evil):
+    result = execute_code(evil)
+    assert not result.success
+
+
+def test_stdin_disabled():
+    result = execute_code("input()")
+    assert not result.success
+
+
+def test_writes_go_to_tempdir():
+    result = execute_code("open('landmine.txt', 'w').write('x'); print(open('landmine.txt').read())")
+    assert result.success and result.stdout == "x\n"
+    assert not os.path.exists("landmine.txt"), "file leaked outside the sandbox tempdir"
+
+
+def test_environment_is_scrubbed():
+    os.environ["FAKE_SECRET_KEY"] = "hunter2"
+    try:
+        result = execute_code("import os; print('FAKE_SECRET_KEY' in os.environ)")
+        assert result.success and result.stdout.strip() == "False"
+    finally:
+        del os.environ["FAKE_SECRET_KEY"]
+
+
+def test_tricky_string_content():
+    # quotes, backslashes and braces must survive being embedded into the runner
+    tricky = '''s = 'it\\'s "quoted" \\\\ {braces} \\n'; print(len(s))'''
+    result = execute_code(tricky)
+    assert result.success
+
+
+def test_fresh_globals():
+    # user code must not see the guard prelude's imports
+    result = execute_code("print('GUARD' in dir())")
+    assert result.success and result.stdout.strip() == "False"

tests/test_optim.py

@@ -0,0 +1,114 @@
+"""
+Test the combined MuonAdamW optimizer (single rank, no process group).
+The AdamW half is checked against torch.optim.AdamW as a reference;
+the Muon half is checked behaviorally (determinism + convergence).
+
+Requires a GPU (the fused kernels are compiled on first use, which
+dominates the runtime of this file).
+
+python -m pytest tests/test_optim.py -v
+"""
+
+import pytest
+import torch
+
+cuda_available = torch.cuda.is_available()
+pytestmark = pytest.mark.skipif(not cuda_available, reason="optimizer tests require CUDA")
+
+if cuda_available:
+    from nanochat.optim import MuonAdamW
+
+DEVICE = "cuda"
+
+# keep shapes small and reuse them across tests so the fused kernels compile once
+ADAMW_SMALL_SHAPE = (48,)        # numel < 1024 path
+ADAMW_LARGE_SHAPE = (1024, 4)    # numel >= 1024 path
+MUON_WIDE_SHAPE = (32, 64)
+MUON_TALL_SHAPE = (64, 32)
+
+
+def make_params_and_groups(seed=1337, offset=0.0, targets=None):
+    """The menagerie: small/large AdamW params, wide/tall Muon stacks."""
+    gen = torch.Generator(device=DEVICE).manual_seed(seed)
+    def rand(shape):
+        base = torch.randn(shape, generator=gen, device=DEVICE) * 0.05
+        return torch.nn.Parameter(base + offset)
+    params = [rand(ADAMW_SMALL_SHAPE), rand(ADAMW_LARGE_SHAPE)]
+    params += [rand(MUON_WIDE_SHAPE) for _ in range(3)]
+    params += [rand(MUON_TALL_SHAPE) for _ in range(2)]
+    groups = [
+        dict(kind="adamw", params=params[0:1], lr=0.02, betas=(0.8, 0.95), eps=1e-10, weight_decay=0.0),
+        dict(kind="adamw", params=params[1:2], lr=0.02, betas=(0.8, 0.96), eps=1e-10, weight_decay=0.0),
+        dict(kind="muon", params=params[2:5], lr=0.02, momentum=0.95, ns_steps=5, beta2=0.9, weight_decay=0.0),
+        dict(kind="muon", params=params[5:7], lr=0.02, momentum=0.95, ns_steps=5, beta2=0.9, weight_decay=0.0),
+    ]
+    return params, groups
+
+
+def test_adamw_matches_torch_reference():
+    """Our fused AdamW must agree with torch.optim.AdamW (both decoupled wd)."""
+    hypers = dict(lr=0.01, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1)
+    p_ours = torch.nn.Parameter(torch.randn(64, 32, device=DEVICE))
+    p_ref = torch.nn.Parameter(p_ours.detach().clone())
+    opt_ours = MuonAdamW([dict(kind="adamw", params=[p_ours], **hypers)])
+    opt_ref = torch.optim.AdamW([p_ref], **hypers)
+    for step in range(10):
+        grad = torch.randn(p_ours.shape, generator=torch.Generator(device=DEVICE).manual_seed(step), device=DEVICE)
+        p_ours.grad = grad.clone()
+        p_ref.grad = grad.clone()
+        opt_ours.step()
+        opt_ref.step()
+    torch.testing.assert_close(p_ours, p_ref, rtol=1e-5, atol=1e-6)
+
+
+def test_determinism():
+    """Two identical runs must produce bitwise identical parameters."""
+    results = []
+    for _ in range(2):
+        params, groups = make_params_and_groups()
+        opt = MuonAdamW(groups)
+        for step in range(5):
+            gen = torch.Generator(device=DEVICE).manual_seed(step)
+            for p in params:
+                p.grad = torch.randn(p.shape, generator=gen, device=DEVICE) * 0.01
+            opt.step()
+        results.append([p.detach().clone() for p in params])
+    for pa, pb in zip(*results):
+        assert torch.equal(pa, pb)
+
+
+def test_convergence():
+    """Optimizing distance-to-target must actually approach the target."""
+    targets_params, _ = make_params_and_groups(seed=999)
+    targets = [p.detach().clone() for p in targets_params]
+    params, groups = make_params_and_groups(seed=999, offset=0.1) # start offset from the targets
+    opt = MuonAdamW(groups)
+    def distances():
+        return [(p.detach() - t).norm().item() for p, t in zip(params, targets)]
+    initial = distances()
+    for _ in range(50):
+        for p, t in zip(params, targets):
+            p.grad = 2 * (p.detach() - t) # gradient of ||p - t||^2
+        opt.step()
+    final = distances()
+    for i, (d0, d1) in enumerate(zip(initial, final)):
+        assert d1 < 0.5 * d0, f"param {i} did not converge: {d0:.4f} -> {d1:.4f}"
+        assert torch.isfinite(params[i]).all()
+
+
+def test_muon_update_is_orthogonalized():
+    """
+    The very first Muon update (zero momentum state) of a full-rank gradient
+    should be (near) semi-orthogonal after the polar iteration: its nonzero
+    singular values land in a band around 1, rather than being spread out.
+    """
+    p = torch.nn.Parameter(torch.zeros(MUON_WIDE_SHAPE, device=DEVICE))
+    group = dict(kind="muon", params=[p], lr=1.0, momentum=0.0, ns_steps=5, beta2=1.0, weight_decay=0.0)
+    opt = MuonAdamW([group])
+    p.grad = torch.randn(p.shape, generator=torch.Generator(device=DEVICE).manual_seed(0), device=DEVICE)
+    opt.step()
+    # with lr=1, wd=0: p_new = -update, so the update is just -p
+    svals = torch.linalg.svdvals(-p.detach().float())
+    # NorMuon variance reduction rescales the magnitude, so normalize by the mean
+    svals = svals / svals.mean()
+    assert svals.max() / svals.min() < 4.0, f"update far from semi-orthogonal: {svals}"

tests/test_tasks.py

@@ -0,0 +1,88 @@
+"""
+Test the Task container machinery: slicing views, mixtures, and the
+HubDataset parquet wrapper (in-memory, no network).
+
+python -m pytest tests/test_tasks.py -v
+"""
+
+import numpy as np
+import pyarrow as pa
+from tasks.common import Task, TaskMixture, HubDataset, render_mc
+
+
+class ToyTask(Task):
+    """A trivial task: example i is just {'i': i, 'tag': tag}."""
+
+    def __init__(self, n=10, tag="a", **kwargs):
+        super().__init__(**kwargs)
+        self.n = n
+        self.tag = tag
+
+    def num_examples(self):
+        return self.n
+
+    def get_example(self, index):
+        return {"i": index, "tag": self.tag}
+
+
+def test_task_full():
+    task = ToyTask(n=10)
+    assert len(task) == 10
+    assert task[0] == {"i": 0, "tag": "a"}
+    assert task[9] == {"i": 9, "tag": "a"}
+
+
+def test_task_slicing():
+    # a view of [5, 10) has 5 examples and maps logical to physical indices
+    task = ToyTask(n=10, start=5, stop=10)
+    assert len(task) == 5
+    assert task[0]["i"] == 5
+    # step slicing uses ceil division for the length
+    task = ToyTask(n=10, start=0, stop=10, step=3) # 0, 3, 6, 9
+    assert len(task) == 4
+    assert [task[i]["i"] for i in range(4)] == [0, 3, 6, 9]
+
+
+def test_mixture_covers_all_examples_deterministically():
+    mixture = TaskMixture([ToyTask(n=3, tag="a"), ToyTask(n=5, tag="b")])
+    assert len(mixture) == 8
+    examples = [mixture[i] for i in range(8)]
+    # every example appears exactly once
+    keys = sorted((ex["tag"], ex["i"]) for ex in examples)
+    assert keys == [("a", 0), ("a", 1), ("a", 2), ("b", 0), ("b", 1), ("b", 2), ("b", 3), ("b", 4)]
+    # the shuffle is deterministic: a second instance yields the same order
+    mixture2 = TaskMixture([ToyTask(n=3, tag="a"), ToyTask(n=5, tag="b")])
+    assert examples == [mixture2[i] for i in range(8)]
+    # and the tasks are actually interleaved, not concatenated
+    assert [ex["tag"] for ex in examples] != ["a"] * 3 + ["b"] * 5
+
+
+def test_mixture_oversampling():
+    # passing a task twice doubles its examples
+    mixture = TaskMixture([ToyTask(n=3), ToyTask(n=3)])
+    assert len(mixture) == 6
+
+
+def test_hub_dataset_rows():
+    table = pa.table({"x": list(range(100)), "y": [str(i) for i in range(100)]})
+    ds = HubDataset(table)
+    assert len(ds) == 100
+    assert ds[7] == {"x": 7, "y": "7"}
+
+
+def test_hub_dataset_shuffle_matches_numpy():
+    # the shuffle must reproduce datasets.Dataset.shuffle(seed) exactly,
+    # which is a np.random.default_rng(seed) permutation
+    table = pa.table({"x": list(range(100))})
+    ds = HubDataset(table).shuffle(seed=42)
+    perm = np.random.default_rng(42).permutation(100)
+    assert [ds[i]["x"] for i in range(100)] == [int(p) for p in perm]
+    # shuffling returns a view; the original order is untouched
+    assert HubDataset(table)[0] == {"x": 0}
+
+
+def test_render_mc_letter_binding():
+    query = render_mc("What is 1+1?", ("A", "B"), ("1", "2"))
+    # the letter must directly follow '=' with no whitespace, so that the
+    # prompt token for "A" matches the assistant's bare "A" response token
+    assert "=A\n" in query and "=B\n" in query

tests/test_tokenizer.py

@@ -0,0 +1,136 @@
+"""
+Test the tokenizer: BPE round-trips, special tokens, and conversation rendering.
+Trains a tiny throwaway tokenizer in-process so the test is hermetic
+(no dependency on ~/.cache/nanochat).
+
+python -m pytest tests/test_tokenizer.py -v
+"""
+
+import pytest
+from nanochat.tokenizer import RustBPETokenizer, SPECIAL_TOKENS
+
+# a small corpus is enough to exercise the BPE machinery
+CORPUS = [
+    "The quick brown fox jumps over the lazy dog.",
+    "hello world, hello tokenizer, hello hello hello",
+    "Numbers like 12345 and unicode like naïve café 你好 🙂 should survive.",
+    "def f(x):\n    return x + 1\n",
+] * 8
+
+
+@pytest.fixture(scope="module")
+def tokenizer():
+    vocab_size = 256 + len(SPECIAL_TOKENS) + 35 # bytes + specials + a few merges
+    return RustBPETokenizer.train_from_iterator(iter(CORPUS), vocab_size)
+
+
+def test_vocab_size(tokenizer):
+    assert tokenizer.get_vocab_size() == 256 + len(SPECIAL_TOKENS) + 35
+
+
+def test_encode_decode_roundtrip(tokenizer):
+    for text in ["hello world", "naïve café 你好 🙂", "unseen tokens: zqxjkv"]:
+        ids = tokenizer.encode(text)
+        assert tokenizer.decode(ids) == text
+
+
+def test_special_tokens(tokenizer):
+    # all special tokens encode to a unique single id
+    ids = [tokenizer.encode_special(t) for t in SPECIAL_TOKENS]
+    assert len(set(ids)) == len(SPECIAL_TOKENS)
+    assert tokenizer.get_bos_token_id() == tokenizer.encode_special("<|bos|>")
+    # specials are NOT special-cased in ordinary text encoding
+    ids = tokenizer.encode("<|bos|>")
+    assert len(ids) > 1, "special token strings in plain text should not collapse to one token"
+
+
+def test_encode_prepend_append(tokenizer):
+    bos = tokenizer.get_bos_token_id()
+    ids = tokenizer.encode("hello", prepend="<|bos|>", append="<|user_end|>")
+    assert ids[0] == bos
+    assert ids[-1] == tokenizer.encode_special("<|user_end|>")
+
+
+def test_encode_batch(tokenizer):
+    texts = ["hello", "world"]
+    ids = tokenizer.encode(texts)
+    assert isinstance(ids, list) and len(ids) == 2
+    assert ids[0] == tokenizer.encode("hello")
+
+
+def test_render_conversation_masks(tokenizer):
+    conversation = {"messages": [
+        {"role": "user", "content": "hi"},
+        {"role": "assistant", "content": "hello!"},
+        {"role": "user", "content": "bye"},
+        {"role": "assistant", "content": "later"},
+    ]}
+    ids, mask = tokenizer.render_conversation(conversation)
+    assert len(ids) == len(mask)
+    # first token is bos and is not supervised
+    assert ids[0] == tokenizer.get_bos_token_id() and mask[0] == 0
+    # supervised tokens are exactly: assistant content + assistant_end tokens
+    assistant_end = tokenizer.encode_special("<|assistant_end|>")
+    supervised_ids = [i for i, m in zip(ids, mask) if m == 1]
+    expected = tokenizer.encode("hello!") + [assistant_end] + tokenizer.encode("later") + [assistant_end]
+    assert supervised_ids == expected
+    # user content tokens are never supervised
+    user_start = tokenizer.encode_special("<|user_start|>")
+    assert all(m == 0 for i, m in zip(ids, mask) if i == user_start)
+
+
+def test_render_conversation_system_message_merged(tokenizer):
+    without_system = {"messages": [
+        {"role": "user", "content": "sys prompt\n\nhi"},
+        {"role": "assistant", "content": "yo"},
+    ]}
+    with_system = {"messages": [
+        {"role": "system", "content": "sys prompt"},
+        {"role": "user", "content": "hi"},
+        {"role": "assistant", "content": "yo"},
+    ]}
+    assert tokenizer.render_conversation(with_system) == tokenizer.render_conversation(without_system)
+
+
+def test_render_conversation_tool_parts(tokenizer):
+    # python tool calls are supervised, python outputs (come from the interpreter) are not
+    conversation = {"messages": [
+        {"role": "user", "content": "add"},
+        {"role": "assistant", "content": [
+            {"type": "text", "text": "sure"},
+            {"type": "python", "text": "1+1"},
+            {"type": "python_output", "text": "2"},
+            {"type": "text", "text": "it is 2"},
+        ]},
+    ]}
+    ids, mask = tokenizer.render_conversation(conversation)
+    python_start = tokenizer.encode_special("<|python_start|>")
+    output_start = tokenizer.encode_special("<|output_start|>")
+    output_end = tokenizer.encode_special("<|output_end|>")
+    # the tool call and its delimiters are supervised
+    assert mask[ids.index(python_start)] == 1
+    # the interpreter output and its delimiters are not
+    start, end = ids.index(output_start), ids.index(output_end)
+    assert all(m == 0 for m in mask[start:end + 1])
+
+
+def test_render_conversation_truncation(tokenizer):
+    conversation = {"messages": [
+        {"role": "user", "content": "hello " * 100},
+        {"role": "assistant", "content": "world " * 100},
+    ]}
+    ids, mask = tokenizer.render_conversation(conversation, max_tokens=32)
+    assert len(ids) == 32 and len(mask) == 32
+
+
+def test_render_for_completion(tokenizer):
+    conversation = {"messages": [
+        {"role": "user", "content": "hi"},
+        {"role": "assistant", "content": "this gets stripped"},
+    ]}
+    ids = tokenizer.render_for_completion(conversation)
+    # ends with assistant_start, primed for a completion
+    assert ids[-1] == tokenizer.encode_special("<|assistant_start|>")
+    # the assistant response itself must not be present
+    stripped = tokenizer.encode("this gets stripped")
+    assert not any(ids[i:i + len(stripped)] == stripped for i in range(len(ids)))