Shared Lab Data

Every module uses the same small set of data, so that you are never learning a new example. Lessons paste in the block they need. This page is the single source for all of them.

All of this data is simulated. Aurora Station is fictional. Every name, log, assignment sheet and image below is generated by the code on this page. Nothing is downloaded: run the blocks and you have the data, the same on every machine.

For authors: copy these blocks into lessons exactly as written. Do not invent different names, word lists, seeds or sizes. If a lesson needs more data, add it in the lesson, say "added in this lesson", and keep what is here unchanged. Every block below was run top to bottom on 9 October 2026 with Python 3.12.15, torch 2.14.1 and transformers 5.19.0, and the outputs quoted on this page come from that run. The timings were measured on the course author's laptop (an Apple silicon Mac), not in Colab. The course standard is one CPU thread (OMP_NUM_THREADS=1), and lesson timings use it; this page also gives the first five-thread run for comparison. Your machine will differ, so measure your own.

What is on this page

Blocks build on each other in order: block 7 needs blocks 2 to 6, and so on. The table in each lesson says which blocks it pastes.

1. Seeds and dates

Every random choice in the course starts from one seed, so your numbers match the course's.

import random
import torch

COURSE_DATE = "2026-10-09"   # versions, model names and prices checked on this day
SEED = 11                    # every generator in the course starts from this seed
SNAPSHOT_SOL = 400           # the corpus is the station record up to sol 399

def set_seed(seed=SEED):
    """Seed Python and PyTorch so a cell gives the same numbers every run."""
    random.seed(seed)
    torch.manual_seed(seed)

set_seed()
print("torch", torch.__version__, "| threads", torch.get_num_threads())

Code walkthrough

• COURSE_DATE records when versions and prices were checked. It is a label, not something code computes with.
• SEED = 11 is the one seed. Every generator on this page is built from it, or from SEED plus a small number, so different datasets do not overlap.
• SNAPSHOT_SOL = 400: a sol is a day on Mars. The station record used as the corpus stops at sol 399; the project's story starts on sol 400 (see 03-the-tinyformer-lab.md).
• set_seed seeds both Python's random module (used to write the documents) and PyTorch (used for weights and batches).
• The print shows your torch version and how many CPU threads PyTorch will use. On the author's laptop it printed torch 2.14.1 | threads 5. Colab's free CPU runtime usually shows fewer.

2. The station's word lists

The corpus is written by small programs that fill templates from these lists. Real pretraining starts from a crawl of the web; this course generates its "crawl", so we know exactly what is in it. That is what makes the deduplication, contamination and mixture experiments in Module 8 checkable.

FIRST = ("Maya Tomas Priya Ines Kofi Lena Omar Hana Dara Ivan Nia Ana Ben Chen Eli Fatima Gus Jo Mo Rina "
         "Sasha Tariq Ulla Viktor Wren Yara Zane Aiko Bruno Cleo Dmitri Esme Farid Greta Hugo Isla Jonas Kira Liam Mira").split()
LAST = ("Chen Okafor Nakamura Silva Haddad Novak Ibarra Petrov Aoki Moreau Larsen Osei Reyes Dubois Kaur "
        "Brandt Ferreira Solberg Almeida Winters Volkov Mensah Castillo Lindqvist Barros Hoffman Duarte Nilsson").split()
ROLES = ["flight engineer", "life-support technician", "systems analyst", "navigator", "quartermaster", "medic",
         "geologist", "hydroponics lead", "comms officer", "reactor technician", "mission planner", "robotics engineer"]
PLACES = ["Habitat A", "Habitat B", "the greenhouse", "the machine shop", "the north airlock", "the south airlock",
          "the observation deck", "the water plant", "the reactor bay", "the cargo hold", "the medical bay", "the comms mast"]
SYSTEMS = ["the oxygen recycler", "the water reclaimer", "the primary array", "the backup array", "the coolant loop",
           "the dust filter", "the pressure seal", "the thermal shield", "the antenna gimbal", "the battery bank",
           "the airlock motor", "the soil monitor"]
VERBS = ["inspected", "repaired", "recalibrated", "replaced", "tested", "cleaned",
         "sealed", "rerouted", "monitored", "documented", "restarted", "balanced"]
PROBLEMS = ["a pressure drop", "a coolant leak", "a stuck valve", "a sensor fault", "a loose cable", "dust in the filter",
            "a slow pump", "a cracked seal", "a drifting clock", "a noisy fan", "a soft short", "an intermittent alarm"]
OUTCOMES = ["the alarm cleared", "readings returned to normal", "the fault repeated after an hour",
            "the part was logged for replacement", "pressure held steady", "the system stayed offline",
            "the fix was temporary", "the reading stabilised"]
UNITS = ["kPa", "degrees", "litres", "hours", "percent", "amps", "volts", "grams"]
CROPS = ["beans", "wheat", "kale", "tomatoes", "potatoes", "peppers", "lettuce", "barley"]
SHIFTS = ["the morning shift", "the evening shift", "the night shift", "the second watch"]
print(len(FIRST) * len(LAST), "possible people,", len(SYSTEMS), "systems,", len(PLACES), "places")

Code walkthrough

• FIRST and LAST are 40 first names and 28 surnames. Any first name can meet any surname, so the corpus has 1,120 possible people, far more than a small model can memorise. It has to learn to copy names from context instead.
• ROLES, PLACES and SYSTEMS are the crew's jobs, the station's twelve areas and its twelve pieces of equipment. The project's cast (03) is drawn from these same lists.
• VERBS, PROBLEMS, OUTCOMES and UNITS fill the maintenance logs. CROPS and SHIFTS fill the greenhouse reports and rotas.
• The print shows 1120 possible people, 12 systems, 12 places.

3. Maintenance logs

The most important document type. A log names a person in full, then refers to them by first name. That repeat-after-introduction pattern is what induction heads latch onto in Module 14, and the log's four facts are the Module 12 task.

def person(rng):
    return f"{rng.choice(FIRST)} {rng.choice(LAST)}"

def number(rng):
    return str(rng.randint(2, 99)) if rng.random() < 0.6 else str(rng.randint(100, 980))

def log_entry(rng):
    """A maintenance log. Returns the text and the four facts a log card holds."""
    name, place, system = person(rng), rng.choice(PLACES), rng.choice(SYSTEMS)
    short, unit, outcome = name.split()[0], rng.choice(UNITS), rng.choice(OUTCOMES)
    lines = [f"Log {rng.randint(1000, 9999)}, sol {rng.randint(1, SNAPSHOT_SOL - 1)}, {rng.choice(SHIFTS)}.",
             f"{name}, {rng.choice(ROLES)}, {rng.choice(VERBS)} {system} in {place} after {rng.choice(PROBLEMS)}.",
             f"{short} recorded {number(rng)} {unit} before the work and {number(rng)} {unit} after it.",
             f"{short} noted that {outcome}."]
    if rng.random() < 0.5:
        lines.append(f"{person(rng)} assisted {short} and signed the entry.")
    card = {"person": name, "system": system, "place": place, "outcome": outcome}
    return "\n".join(lines), card

Code walkthrough

• person(rng) picks a first name and a surname. Every generator takes rng, a seeded random.Random, so the same seed always writes the same documents.
• number(rng) gives a two-digit number 60% of the time and a three-digit one otherwise, so numbers of different lengths appear for the tokeniser lessons.
• log_entry picks a person, a place, a system, a unit and an outcome, then writes four lines: a header with a log number, a sol and a shift; who did what to which system, where, after which problem; two readings; and the outcome.
• short is the first name alone. It appears three more times after the full name.
• Half the logs get a fifth line naming a second person who signed the entry. That second name is a distractor: it is the wrong answer to "who did the work".
• The function returns the text and card, the four facts Module 12 trains the model to extract.

4. Assignment sheets with questions

The corpus's built-in reading test. A sheet lists who works where, then asks about one of them. The answer is always earlier in the same document, so answering means looking back and copying.

def fact_block(rng):
    """An assignment sheet, then questions whose answers sit earlier in the same document."""
    names, k = [], rng.randint(3, 6)
    while len(names) < k:                        # no name twice, so every question has one answer
        name = person(rng)
        if name not in names:
            names.append(name)
    rows = [(p, rng.choice(PLACES), rng.choice(SYSTEMS)) for p in names]
    lines = [f"Assignment sheet for sol {rng.randint(1, SNAPSHOT_SOL - 1)}:"]
    lines += [f"{p} works in {place} on {system}." for p, place, system in rows]
    for p, place, system in rng.sample(rows, k=2):
        if rng.random() < 0.5:
            lines += [f"Question: where does {p} work?", f"Answer: {p} works in {place}."]
        else:
            lines += [f"Question: which system does {p} look after?", f"Answer: {p} looks after {system}."]
    return "\n".join(lines)

Code walkthrough

• k people, between three and six, are drawn. The while loop skips a name already on the sheet, so no question can have two right answers.
• Each row gives a person a place and a system: "Kira Ibarra works in the greenhouse on the water reclaimer."
• rng.sample(rows, k=2) picks two different people to ask about. A coin flip decides whether the question is about the place or the system.
• The answer line repeats the person's full name and then the fact. A model that has learnt to look back gets it right; one that only knows which words are common does not.

5. Dialogues, greenhouse reports and procedures

Three more document types, so the corpus is a mixture of sources like a real one. Module 8 changes the mixture weights on purpose.

def dialogue(rng):
    a, b = rng.choice(FIRST), rng.choice(FIRST)
    system, problem = rng.choice(SYSTEMS), rng.choice(PROBLEMS)
    return "\n".join([f"{a}: {b}, {system} is showing {problem}.",
                      f"{b}: How long has it been like that, {a}?",
                      f"{a}: About {rng.randint(2, 40)} hours. The reading sits at {number(rng)} {rng.choice(UNITS)}.",
                      f"{b}: I will check it during {rng.choice(SHIFTS)}.",
                      f"{a}: Thanks, {b}. I will log it."])

def greenhouse(rng):
    crop = rng.choice(CROPS)
    return "\n".join([f"Greenhouse report, sol {rng.randint(1, SNAPSHOT_SOL - 1)}.",
                      f"The {crop} bed grew {number(rng)} millimetres this week.",
                      f"Water use was {number(rng)} litres, {rng.choice(['above', 'below', 'close to'])} the plan.",
                      f"{person(rng)} will thin the {crop} on {rng.choice(SHIFTS)}."])

STEPS = ["close the isolation valve", "remove the access panel", "check the seal for cracks",
         "replace the filter cartridge", "torque the bolts to 12 newton metres", "restore pressure slowly",
         "watch the gauge for five minutes", "log the reading in the station record"]

def procedure(rng):
    lines = [f"Procedure: servicing {rng.choice(SYSTEMS)}."]
    lines += [f"Step {i}: {s}." for i, s in enumerate(rng.sample(STEPS, k=rng.randint(3, 5)), 1)]
    lines.append(f"If {rng.choice(PROBLEMS)} appears, stop and call {person(rng)}.")
    return "\n".join(lines)

GENERATORS = {"log": (lambda rng: log_entry(rng)[0], 0.30), "facts": (fact_block, 0.30),
              "dialogue": (dialogue, 0.15), "greenhouse": (greenhouse, 0.10), "procedure": (procedure, 0.15)}

Code walkthrough

• dialogue is a five-line exchange between two crew members, by first name. Each name is used three times, another chance for the model to learn copying.
• greenhouse writes a short crop report with a growth figure, water use and who will thin the bed.
• procedure picks three to five of the eight STEPS in a random order and numbers them. The numbered steps teach the model to count in order.
• GENERATORS names the five sources and their mixture weights. The weights add to 1.0. log uses a small lambda to keep only the text from log_entry, which also returns a card.

6. Building the corpus

import time
from collections import Counter

def build_corpus(n_docs=6000, seed=SEED, valid_fraction=0.05):
    """Sample n_docs documents from the mixture. Returns (train, valid) lists of (kind, text)."""
    rng = random.Random(seed)
    kinds = list(GENERATORS)
    weights = [GENERATORS[k][1] for k in kinds]
    docs = []
    for _ in range(n_docs):
        kind = rng.choices(kinds, weights)[0]
        docs.append((kind, GENERATORS[kind][0](rng)))
    split = int(n_docs * (1 - valid_fraction))
    return docs[:split], docs[split:]

t0 = time.perf_counter()
TRAIN_DOCS, VALID_DOCS = build_corpus()
print(f"{len(TRAIN_DOCS):,} train docs, {len(VALID_DOCS):,} valid docs,",
      f"{sum(len(t) for _, t in TRAIN_DOCS):,} train characters, {time.perf_counter() - t0:.2f} s")
print(Counter(k for k, _ in TRAIN_DOCS).most_common())
print(TRAIN_DOCS[0][1])

Code walkthrough

• build_corpus makes one seeded random.Random and draws 6,000 documents. For each it picks a source with rng.choices(kinds, weights), then calls that source's generator.
• Every document is kept with its kind. Real corpora lose this label easily; keeping it from the start is what lets Module 8 report the mixture.
• The last 5% become the validation set. The model never trains on them.
• Output on the author's laptop: 5,700 train docs, 300 valid docs, 1,795,697 train characters, 0.02 s, and the mix facts 1738, log 1696, procedure 891, dialogue 848, greenhouse 527. The first document is an assignment sheet for sol 379 that begins "Bruno Ferreira works in the water plant on the coolant loop."

7. The character tokeniser

The model reads numbers, not letters. The simplest honest choice is one token per character: no training, no hidden rules, and every lesson can see exactly what the model sees. Module 8 trains a byte pair encoding tokeniser and compares.

SPECIALS = ["<pad>", "<doc>", "<unk>"]   # padding, document boundary, unseen character
CHARS = sorted(set("".join(t for _, t in TRAIN_DOCS)))
VOCAB = SPECIALS + CHARS
STOI = {s: i for i, s in enumerate(VOCAB)}
PAD, DOC, UNK = 0, 1, 2

def encode(text):
    return [STOI.get(c, UNK) for c in text]

def decode(ids):
    return "".join(VOCAB[i] if i >= len(SPECIALS) else "" for i in ids)

def to_stream(docs):
    """One long tensor of token ids, each document preceded by <doc>."""
    ids = []
    for _, text in docs:
        ids += [DOC] + encode(text)
    return torch.tensor(ids, dtype=torch.long)

TRAIN_IDS, VALID_IDS = to_stream(TRAIN_DOCS), to_stream(VALID_DOCS)
print("vocab", len(VOCAB), "| train tokens", f"{len(TRAIN_IDS):,}", "| valid tokens", f"{len(VALID_IDS):,}")
print(encode("Maya"), decode(encode("Maya")))

Code walkthrough

• SPECIALS are three tokens that are not characters: <pad> fills short sequences, <doc> marks where one document ends and the next begins, <unk> stands for a character never seen in training.
• CHARS is every distinct character in the training text, sorted so the order is always the same. VOCAB puts the specials first, so their ids are 0, 1 and 2.
• STOI ("string to index") maps each token to its id. encode turns text into ids; decode turns ids back into text and drops the specials.
• to_stream joins all documents into one long tensor, each starting with <doc>. Training reads windows from this stream.
• Output: vocab 69 | train tokens 1,801,397 | valid tokens 90,984, then [32, 45, 68, 45] Maya. The vocabulary is 66 characters plus 3 specials, and ModelConfig below uses exactly that size, so there are no unused rows in the embedding table. The corpus never uses {, " or }, so they are not in the vocabulary: anything a model must write, such as Module 12's log cards, uses field: value lines, never JSON.

8. Batches

def get_batch(stream, batch_size=32, context=128, generator=None):
    """Random windows from a token stream: inputs, and the same windows shifted one to the left."""
    starts = torch.randint(0, len(stream) - context - 1, (batch_size,), generator=generator)
    x = torch.stack([stream[s:s + context] for s in starts])
    y = torch.stack([stream[s + 1:s + context + 1] for s in starts])
    return x, y

def fixed_batches(stream, n=8, batch_size=32, context=128, seed=SEED):
    """The same evaluation windows every time, so validation loss is comparable across runs."""
    g = torch.Generator().manual_seed(seed)
    return [get_batch(stream, batch_size, context, g) for _ in range(n)]

VALID_BATCHES = fixed_batches(VALID_IDS)
x, y = VALID_BATCHES[0]
print(x.shape, y.shape, repr(decode(x[0, :40].tolist())))

Code walkthrough

• get_batch picks batch_size random start points in the stream and cuts a window of context tokens from each. y is the same window moved one token to the right: at every position, the target is the next token. That is the whole of language-model training data.
• Passing a generator makes the random starts reproducible without disturbing other random numbers.
• fixed_batches draws eight batches once, from their own seeded generator. VALID_BATCHES is the course's one validation set: every lesson that quotes a validation loss means the average over these eight batches, 32 windows of 128 tokens each.
• Output: torch.Size([32, 128]) torch.Size([32, 128]) and the start of the first window, 'airlock motor.\nBruno Hoffman works in th'. Windows start anywhere, even mid-document.

9. Synthetic tasks

Language mixes many skills. A synthetic task isolates one, and its ceiling is known: 100% is possible, so a model that stalls has failed, not hit a hard task.

from typing import NamedTuple

class TaskData(NamedTuple):
    name: str
    vocab_size: int
    inputs: torch.Tensor    # (n, T) token ids
    targets: torch.Tensor   # (n, T) the next token at every position
    mask: torch.Tensor      # (n, T) 1 where the answer is scored, 0 elsewhere

def _shift(name, vocab, seq, scored_from):
    inputs, targets = seq[:, :-1], seq[:, 1:]
    mask = torch.zeros_like(inputs)
    mask[:, scored_from:] = 1
    return TaskData(name, vocab, inputs, targets, mask)

def make_task(name, n=4096, length=16, vocab=32, seed=SEED):
    """copy, reverse or induction. Token 0 is the separator; content tokens are 1 to vocab - 1."""
    g = torch.Generator().manual_seed(seed)
    body = torch.randint(1, vocab, (n, length), generator=g)
    sep = torch.zeros(n, 1, dtype=torch.long)
    if name == "copy":
        seq = torch.cat([body, sep, body], dim=1)
    elif name == "reverse":
        seq = torch.cat([body, sep, body.flip(1)], dim=1)
    elif name == "induction":                       # random tokens, then the same random tokens again
        seq = torch.cat([body, body], dim=1)
    return _shift(name, vocab, seq, scored_from=length)
TASKS = {name: make_task(name) for name in ["copy", "reverse", "induction"]}
for t in TASKS.values():
    print(t.name, tuple(t.inputs.shape), "scored positions per row:", int(t.mask[0].sum()))

Code walkthrough

• TaskData is a named bundle: the task's name, its vocabulary size, inputs, targets and a mask.
• _shift makes inputs and next-token targets from a full sequence, then sets the mask to 1 only from position scored_from on. The first half is random and cannot be predicted, so it is never scored.
• make_task draws body, 4,096 rows of 16 random tokens from 1 to 31. copy is body, a separator (token 0), body again. reverse is body, separator, body backwards. induction is body twice with no separator: the model must notice that the sequence is repeating, the behaviour that induction heads implement.
• A warning about induction: every copy sits exactly 16 positions back, in every row. A model can solve it with a fixed-offset shortcut ("read the token 15 back") and never learn induction. Module 14, Lesson 1 measured this: a model that passes this task falls to chance when the gap changes. Module 14, Lesson 2 adds make_varied, which gives each row its own gap, and the Module 14 lab saves it in tinyformer/analysis.py as the package's induction task. TASKS["induction"] stays here unchanged, as the course's example of a shortcut.
• Output: copy and reverse are (4096, 32) with 16 scored positions per row; induction is (4096, 31) with 15, because the first token of the repeat cannot be predicted. Chance accuracy on a scored position is 1 in 31, about 3.2%.

10. Lookup questions with expected answers

A held-out test of reading from context, built from fresh assignment sheets.

def lookup_questions(n=200, seed=SEED + 1):
    """Held-out in-context lookups: a prompt that ends at 'Answer:' and the exact expected answer."""
    rng = random.Random(seed)
    items = []
    while len(items) < n:
        lines = fact_block(rng).split("\n")
        q = next(i for i, line in enumerate(lines) if line.startswith("Question:"))
        prompt = "\n".join(lines[:q + 1]) + "\nAnswer:"
        answer = lines[q + 1][len("Answer:"):]                  # for example " Ana Hoffman works in Habitat B."
        target = answer.split(" works in ")[-1].split(" looks after ")[-1].rstrip(".")
        items.append({"prompt": prompt, "answer": answer, "target": target})
    return items

LOOKUP_EVAL = lookup_questions()
print(len(LOOKUP_EVAL), "questions")
print(LOOKUP_EVAL[0]["prompt"])
print("expected:", repr(LOOKUP_EVAL[0]["answer"]), "| target:", LOOKUP_EVAL[0]["target"])

Code walkthrough

• lookup_questions uses its own seed, SEED + 1, so these sheets are written by a different random stream from the training corpus. Module 8 checks how many of them still overlap with training text by chance.
• For each new sheet it finds the first Question: line, keeps everything up to it, and ends the prompt with Answer:. The model must continue from there.
• answer is the full expected line. target is just the fact (a place or a system), which is what a lesson scores: did the model's reply contain it?
• Output: 200 questions, then a six-person sheet for sol 174 asking which system does Greta Mensah look after?, with expected answer ' Greta Mensah looks after the backup array.' and target the backup array.
• A limit to know before you score base on these. A prompt runs 233 to 457 characters (351 on average), but base reads only the last 128. On the first 50 questions, the sheet line that holds the answer sits wholly inside those last 128 characters for only 11 (45 of all 200). For the rest, the fact has been cut off before the model sees it. So a low lookup score from base is partly truncation, not only a missing skill. Later lessons that want to test reading from context either use a longer context or shorter prompts, and say which.

11. Log cards: the Module 12 task

Maya Lindqvist wants every maintenance log turned into a four-field card (see 03). These are the supervised pairs.

def card_text(card):
    return (f"person: {card['person']}\nsystem: {card['system']}\n"
            f"place: {card['place']}\noutcome: {card['outcome']}")

def log_card_set(n, seed):
    """Supervised pairs for Module 12: a maintenance log in, a four-field card out."""
    rng = random.Random(seed)
    pairs = []
    for _ in range(n):
        text, card = log_entry(rng)
        pairs.append({"prompt": text + "\nCard:\n", "completion": card_text(card), "card": card})
    return pairs

CARD_TRAIN = log_card_set(1000, seed=SEED + 2)
CARD_EVAL = log_card_set(200, seed=SEED + 3)
print(len(CARD_TRAIN), "train pairs,", len(CARD_EVAL), "eval pairs")
print(CARD_EVAL[0]["prompt"] + CARD_EVAL[0]["completion"])

Code walkthrough

• card_text writes a card as four field: value lines, always in the same order. The format is fixed so a check can compare it exactly.
• log_card_set writes n fresh logs with its own seed. Each pair has a prompt (the log, then Card:), a completion (the card the model should write) and the card as a dictionary for scoring field by field.
• Training and evaluation use different seeds (SEED + 2 and SEED + 3), so no evaluation log is a training log.
• A log and its card together run about 305 to 440 characters, so none fits base's 128-character context. Module 12 trains and scores cards with a 448-character window, CARD_WINDOW = 448.
• Output: 1000 train pairs, 200 eval pairs, then log 5190 from sol 150, in which Omar Winters, mission planner, rerouted the airlock motor in the comms mast and Sasha Osei signed the entry. The card names Omar Winters, not Sasha Osei: the distractor at work.

12. Drawing one shape

Module 13 needs pictures. These are tiny generated images, small enough to train a vision transformer on a CPU in seconds.

SHAPES = ["square", "circle", "triangle"]
COLOURS = {"red": (0.85, 0.2, 0.2), "green": (0.2, 0.7, 0.3), "blue": (0.2, 0.35, 0.85), "yellow": (0.9, 0.8, 0.2)}
WORDS = {1: "one", 2: "two", 3: "three"}

def draw(img, shape, rgb, top, left, size):
    """Paint one shape into a (3, 32, 32) image, inside a size-by-size box at (top, left)."""
    r = torch.arange(size).view(-1, 1).float()
    c = torch.arange(size).view(1, -1).float()
    if shape == "square":
        m = torch.ones(size, size, dtype=torch.bool)
    elif shape == "circle":
        h = (size - 1) / 2
        m = (r - h) ** 2 + (c - h) ** 2 <= (size / 2) ** 2
    else:                                            # triangle: apex at the top, base at the bottom
        m = (c - (size - 1) / 2).abs() <= r / 2
    for ch in range(3):
        img[ch, top:top + size, left:left + size][m] = rgb[ch]

Code walkthrough

• Three shapes, four colours as red, green and blue amounts between 0 and 1, and number words for captions.
• draw builds a true/false mask m the size of the shape's box. r is the row number and c the column number of every pixel in the box.
• A square fills the whole box. A circle keeps pixels within half the box's width of its centre. A triangle keeps pixels whose distance from the middle column is at most half the row number, so it is one pixel wide at the top and full width at the bottom.
• The last loop paints the masked pixels in each of the three colour channels.

13. The shapes dataset

def make_shapes(n=2000, seed=SEED):
    """n images, each holding 1 to 3 copies of one shape in one colour, one per 16x16 quarter."""
    g = torch.Generator().manual_seed(seed)
    images = torch.zeros(n, 3, 32, 32)
    shape_y, colour_y, count_y, captions = [], [], [], []
    for i in range(n):
        s = int(torch.randint(3, (1,), generator=g))
        c = int(torch.randint(4, (1,), generator=g))
        k = int(torch.randint(1, 4, (1,), generator=g))
        name = list(COLOURS)[c]
        for q in torch.randperm(4, generator=g)[:k].tolist():
            size = int(torch.randint(7, 13, (1,), generator=g))
            top = (q // 2) * 16 + int(torch.randint(0, 17 - size, (1,), generator=g))
            left = (q % 2) * 16 + int(torch.randint(0, 17 - size, (1,), generator=g))
            draw(images[i], SHAPES[s], COLOURS[name], top, left, size)
        shape_y.append(s); colour_y.append(c); count_y.append(k)
        captions.append(f"{WORDS[k]} {name} {SHAPES[s]}{'s' if k > 1 else ''}")
    return images, torch.tensor(shape_y), torch.tensor(colour_y), torch.tensor(count_y), captions

t0 = time.perf_counter()
IMAGES, SHAPE_Y, COLOUR_Y, COUNT_Y, CAPTIONS = make_shapes()
print(IMAGES.shape, f"{time.perf_counter() - t0:.2f} s", CAPTIONS[:3])
print("shape counts", torch.bincount(SHAPE_Y).tolist(), "| object counts", torch.bincount(COUNT_Y)[1:].tolist())

Code walkthrough

• Each image is 32 by 32 pixels, black, split into four 16 by 16 quarters. It holds one, two or three copies of the same shape in the same colour, at most one per quarter, so shapes never overlap.
• s, c and k are the shape, colour and count, drawn from a seeded generator. torch.randperm(4)[:k] picks which quarters to use.
• Each copy gets its own size (7 to 12 pixels) and a random position inside its quarter.
• Four labels come out: SHAPE_Y, COLOUR_Y, COUNT_Y and a caption such as "two red circles". Module 13 trains on shape, then on count, and uses the captions for CLIP and the captioner.
• Output: torch.Size([2000, 3, 32, 32]) in 0.06 s, first captions 'one yellow square', 'one blue triangle', 'one green circle', shapes balanced at 686, 666 and 648, and counts of one, two and three at 638, 697 and 665.

14. Looking at an image without a plotting library

def show(i):
    """Print image i as text: '#' where any colour is present, '.' where it is black."""
    lit = IMAGES[i].amax(dim=0) > 0
    print(CAPTIONS[i], "| labels:", SHAPES[SHAPE_Y[i]], list(COLOURS)[COLOUR_Y[i]], int(COUNT_Y[i]))
    for row in lit[::2]:                                    # every second row keeps it short
        print("".join("#" if v else "." for v in row.tolist()))

# image 3 holds three shapes, image 1 holds one
show(3)
show(1)

Code walkthrough

• IMAGES[i].amax(dim=0) takes the brightest of the three colour channels at each pixel; anything above zero is part of a shape.
• The loop prints every second row as a line of # and ., enough to check a shape and a count by eye. Matplotlib works too (plt.imshow(IMAGES[i].permute(1, 2, 0))); this version needs nothing.
• The caption line also prints the three numeric labels decoded back to words, so you can check they agree with the caption.
• Output: image 3 is three red circles (16 lines of 32 characters showing three round blobs in three quarters); image 1 is one blue triangle.

15. The reference model: configuration and rotary positions

Modules 2 to 5 build the transformer piece by piece, with every switch explained. From Module 6 on, lessons need a working model without rebuilding it, so this page keeps the finished, compact reference: the same design as the lab's tinyformer/model.py, cut to the switches later lessons use. Blocks 15 to 19 are pasted together.

import math
from dataclasses import dataclass
import torch.nn as nn
import torch.nn.functional as F

@dataclass
class ModelConfig:
    vocab_size: int = 69
    context: int = 128
    d_model: int = 128
    n_layers: int = 4
    n_heads: int = 4
    n_kv_heads: int = 4          # fewer than n_heads gives grouped-query attention; 1 gives multi-query
    d_ff: int = 344              # hidden width of the feed-forward network
    position: str = "rope"       # "rope", "learned" or "none"
    norm: str = "rms"            # "rms" or "layer"
    ffn: str = "swiglu"          # "swiglu" or "gelu"
    tie_embeddings: bool = True

def make_norm(cfg):
    return nn.RMSNorm(cfg.d_model) if cfg.norm == "rms" else nn.LayerNorm(cfg.d_model)

def rope(x, base=10000.0):
    """Rotate pairs of features by an angle that grows with position. x: (B, heads, T, head_dim)."""
    T, d = x.shape[-2], x.shape[-1]
    freqs = base ** (-torch.arange(0, d, 2, device=x.device).float() / d)
    angles = torch.arange(T, device=x.device).float()[:, None] * freqs[None, :]
    cos, sin = angles.cos().to(x.dtype), angles.sin().to(x.dtype)   # keep 16-bit inputs 16-bit
    x1, x2 = x[..., 0::2], x[..., 1::2]
    return torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).flatten(-2)

Code walkthrough

• ModelConfig holds every design choice as a named switch. The defaults are the course's base model: 4 layers, 4 heads, width 128, context 128 tokens, rotary positions, RMSNorm, SwiGLU, tied embeddings.
• n_kv_heads sets how many key/value heads there are. Equal to n_heads is ordinary multi-head attention; smaller shares each key/value head across several query heads (Module 2, Lesson 5).
• d_ff = 344: SwiGLU uses three weight matrices where GELU uses two, so its hidden width is two thirds of the usual 4 x 128 = 512, rounded up to a multiple of 8. That keeps the two the same size: 132,096 weights per layer against 131,072. When you switch ffn to "gelu", set d_ff=512 too (Module 3, Lesson 1 derives this).
• make_norm returns RMSNorm or LayerNorm according to the switch.
• rope rotates each pair of features in a query or key by an angle proportional to the token's position, with a different speed for each pair. Module 4, Lesson 2 builds it from a picture. The angles are computed in float32 and then cast to the input's type, so a bfloat16 or float16 model stays in that type.

16. Attention with grouped key/value heads

class Attention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg, self.hd = cfg, cfg.d_model // cfg.n_heads
        self.q = nn.Linear(cfg.d_model, cfg.n_heads * self.hd, bias=False)
        self.kv = nn.Linear(cfg.d_model, 2 * cfg.n_kv_heads * self.hd, bias=False)
        self.out = nn.Linear(cfg.n_heads * self.hd, cfg.d_model, bias=False)

    def forward(self, x):
        B, T, _ = x.shape
        q = self.q(x).view(B, T, self.cfg.n_heads, self.hd).transpose(1, 2)
        k, v = self.kv(x).view(B, T, 2, self.cfg.n_kv_heads, self.hd).permute(2, 0, 3, 1, 4)
        if self.cfg.position == "rope":
            q, k = rope(q), rope(k)
        repeat = self.cfg.n_heads // self.cfg.n_kv_heads          # share each key/value head
        k, v = k.repeat_interleave(repeat, dim=1), v.repeat_interleave(repeat, dim=1)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.out(y.transpose(1, 2).reshape(B, T, -1))

Code walkthrough

• self.hd is the width of one head: 128 / 4 = 32.
• self.q makes the queries for every head. self.kv makes keys and values together, for only n_kv_heads heads. self.out mixes the heads' results back to the model width.
• In forward, view and transpose reshape the queries to (batch, heads, tokens, head width). The keys and values are split and reshaped the same way with one permute.
• With position == "rope", queries and keys are rotated by position before they are compared.
• repeat_interleave copies each key/value head so every query head has a partner. With n_kv_heads == n_heads it copies nothing.
• F.scaled_dot_product_attention(..., is_causal=True) is PyTorch's built-in attention with the causal mask: each token sees itself and earlier tokens only. Module 2 writes the same thing by hand and checks they agree.

17. Feed-forward network, block and model

class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.up = nn.Linear(cfg.d_model, cfg.d_ff * (2 if cfg.ffn == "swiglu" else 1), bias=False)
        self.down = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)

    def forward(self, x):
        if self.cfg.ffn == "swiglu":
            gate, value = self.up(x).chunk(2, dim=-1)
            return self.down(F.silu(gate) * value)
        return self.down(F.gelu(self.up(x)))

class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.norm1, self.attn = make_norm(cfg), Attention(cfg)
        self.norm2, self.ffn = make_norm(cfg), FeedForward(cfg)

    def forward(self, x):                       # pre-norm: normalise, transform, add back to the stream
        x = x + self.attn(self.norm1(x))
        return x + self.ffn(self.norm2(x))

Code walkthrough

• FeedForward widens each token's vector, applies a non-linearity, and narrows it back. With SwiGLU, up produces two halves: gate decides how much of value passes through (F.silu(gate) * value). With GELU there is one half and a plain activation.
• Block is one transformer layer: attention, then the feed-forward network, each wrapped as "normalise, transform, add the result back". The running sum x is the residual stream (Module 3, Lesson 2).
• This is the pre-norm arrangement: the normalisation happens before each sub-layer, not after. Module 3, Lesson 3 measures why that matters.

18. The whole model

class TinyFormer(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.pos = nn.Embedding(cfg.context, cfg.d_model) if cfg.position == "learned" else None
        self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layers))
        self.norm = make_norm(cfg)
        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        if cfg.tie_embeddings:
            self.head.weight = self.embed.weight
        for name, p in self.named_parameters():
            if p.dim() == 2:
                nn.init.normal_(p, 0.0, 0.02)

    def forward(self, ids, targets=None):
        x = self.embed(ids)
        if self.pos is not None:
            x = x + self.pos(torch.arange(ids.shape[1], device=ids.device))
        for block in self.blocks:
            x = block(x)
        logits = self.head(self.norm(x))
        if targets is None:
            return logits
        return logits, F.cross_entropy(logits.flatten(0, 1), targets.flatten())

    def n_params(self):
        return sum(p.numel() for p in self.parameters())

Code walkthrough

• embed turns each token id into a vector of width d_model. pos is a second table, one vector per position, used only when position == "learned".
• blocks is the stack of layers. norm is the final normalisation and head turns each vector into one score per vocabulary entry, the logits.
• With tie_embeddings, the output head reuses the embedding table's weights, so the model has one table, not two.
• The loop at the end gives every weight matrix small random starting values (standard deviation 0.02). Module 5, Lesson 2 measures what happens with other choices.
• forward returns the logits alone, or the logits and the average cross-entropy loss when targets are given. n_params counts the weights.

19. Training, evaluating and the base recipe

This replaces the old course's shipped model. Instead of downloading a trained model, every lesson that needs one trains base from the same seed in about one to two minutes on the author's laptop with one thread.

def train(model, steps=600, lr=3e-3, batch_size=32, warmup=50, seed=SEED, log_every=200):
    """AdamW, linear warm-up then cosine decay, gradient clipping at 1.0. Returns the final validation loss."""
    g = torch.Generator().manual_seed(seed)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=0.1)
    schedule = lambda s: min(1.0, (s + 1) / warmup) * 0.5 * (1 + math.cos(math.pi * s / steps))
    sched = torch.optim.lr_scheduler.LambdaLR(opt, schedule)
    model.train()
    for step in range(steps):
        x, y = get_batch(TRAIN_IDS, batch_size, model.cfg.context, g)
        _, loss = model(x, y)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step(); sched.step()
        if (step + 1) % log_every == 0:
            print(f"step {step + 1:4d}  train loss {loss.item():.3f}  valid loss {evaluate(model):.3f}")
    return evaluate(model)

@torch.no_grad()
def evaluate(model, batches=VALID_BATCHES):
    model.eval()
    loss = sum(model(x, y)[1].item() for x, y in batches) / len(batches)
    model.train()
    return loss

Code walkthrough

• train uses AdamW, the standard optimiser for transformers, with the settings Module 6 explains: betas 0.9 and 0.95, weight decay 0.1.
• schedule raises the learning rate from almost zero over the first 50 steps (warm-up), then lowers it along a cosine curve to zero at the last step.
• Each step draws a batch from its own seeded generator, computes the loss, clears old gradients, back-propagates, clips the gradient norm to 1.0 and updates the weights.
• Every log_every steps it prints the training loss of the current batch and the validation loss.
• evaluate averages the loss over the eight fixed VALID_BATCHES with gradients switched off (@torch.no_grad()), and puts the model back into training mode after.

@torch.no_grad()
def generate(model, prompt, max_new=80, temperature=0.0):
    """Greedy when temperature is 0, sampling otherwise. Stops at the end of a document."""
    model.eval()
    ids = torch.tensor([[DOC] + encode(prompt)])
    for _ in range(max_new):
        logits = model(ids[:, -model.cfg.context:])[0, -1]
        nxt = logits.argmax() if temperature == 0 else torch.multinomial(F.softmax(logits / temperature, -1), 1)[0]
        if nxt.item() == DOC:
            break
        ids = torch.cat([ids, nxt.view(1, 1)], dim=1)
    return decode(ids[0, 1:].tolist())

set_seed()
base = TinyFormer(ModelConfig(vocab_size=len(VOCAB)))
print(f"{base.n_params():,} parameters")
t0 = time.perf_counter()
print(f"final valid loss {train(base):.3f} in {time.perf_counter() - t0:.0f} s")
print(generate(base, "Log 4821, sol 212, the night shift.\n", max_new=120))
hits = 0
for q in LOOKUP_EVAL[:50]:
    reply = generate(base, q["prompt"], max_new=60)[len(q["prompt"]):].split("\n")[0]
    hits += q["target"] in reply
print(f"in-context lookup: {hits} of 50 correct")

Code walkthrough

• generate starts from <doc> and the encoded prompt, and adds one token at a time. Each step runs the model on the last context tokens and takes the scores for the final position.
• With temperature == 0 it takes the highest-scoring token (greedy decoding). Otherwise it divides the scores by the temperature, turns them into probabilities and samples. It stops if the model starts a new document.
• This simple loop recomputes the whole window every step. Module 11 adds a KV cache and measures the difference.
• The base recipe: set_seed() first, so the starting weights are the course's; ModelConfig(vocab_size=len(VOCAB)) takes every default and the tokeniser's real vocabulary size; train(base) runs the 600 steps. Output on the author's laptop with one thread (13 October 2026): 800,512 parameters, then validation loss 0.372 at step 200, 0.274 at step 400 and 0.256 at step 600, in 66 s. Across the course's one-thread runs, training took about 50 to 110 s, depending on what else the machine was doing; with five threads it took 25 to 26 s, and the step 200 and 400 losses differed in the last digit (0.370, 0.276). The final 0.256 matched. The loss is in nats per character. It is low because most of each document is template text that is easy to predict; the names and numbers are the hard part.
• The first print continues a log header. Do not expect the same words as the course: the names and numbers in the sample change with the thread count, even when the loss matches. A five-thread run gave "Mira Nakamura, robotics engineer, replaced ..." and the next line began "Mo recorded"; a one-thread run gave "Maya Hoffman, robotics engineer, replaced the primary array in the cargo hold after a sensor fault." and then "Tomas recorded". What stays the same is the pattern: the grammar is right, but the first name in the next line is not the person's own. The model has not yet learnt to copy a name.
• The loop scores the first 50 lookup questions: does the reply contain the target? Output with one thread: in-context lookup: 3 of 50 correct (a five-thread run gave 4 of 50). Two things keep it low. First, truncation: the prompts are longer than base's 128-character context, so generate sees only the last 128 characters, and for 39 of these 50 questions the line holding the answer is not wholly inside them (block 10). Second, base has not learnt to copy a fact from its context; Module 9 found that no model in its grid, base included, wrote a full answer line exactly right (0 of 200). Module 10, Lesson 4 tests reading from context with its own probes that fit the window, so the two causes can be told apart. This gap is the thread that runs through Modules 8, 10 and 14: what a model needs before it can read from its own context.

Timings for this page (author's laptop, not Colab)

Five-thread runs: 9 October 2026; one-thread run: 13 October 2026; both on an Apple silicon Mac. A free Colab CPU runtime has fewer cores; budget up to about two minutes for base. Numbers that come from seeded code (losses, counts, accuracies) should match yours to the digits shown when you run on the same versions; on a different machine or thread count the last digit of a loss can differ, a lookup count can move by one, and generated text can change.