Delete cis6270

cis6270/__init__.py

@@ -1,20 +0,0 @@
-"""Shared helpers for the CIS 6270 course code.
-
-Nothing in this package is the subject of a lecture. The data, the small networks, the training
-loop and the sample metrics live here so that each lecture file holds only the mathematics of its
-own framework. Read this package once; after that every lecture file stands on its own.
-"""
-
-import warnings
-
-# torch 2.14 warns about `torch.jit.script` from inside its own forward-mode autograd machinery,
-# the first time `torch.func.jvp` runs. Chapter 6 uses that function, and nothing here is
-# deprecated, so the message is filtered and the reader never sees it.
-warnings.filterwarnings(
-    "ignore", message=r"`torch\.jit\.script` is deprecated", category=FutureWarning,
-)
-
-from cis6270 import data, metrics, nets, runner
-
-__all__ = ["data", "metrics", "nets", "runner"]
-__version__ = "1.0.0"

cis6270/data.py

@@ -1,317 +0,0 @@
-"""Datasets for the course.
-
-Everything here runs offline in under a second, so a fresh clone and the test suite need no
-network. MNIST and ESM-2 embeddings download on first use and are reached only through the
-``--dataset`` flag of the scripts that offer them.
-
-The discrete lectures all use the same synthetic DNA, so a loss that improves in Lecture 4 and a
-loss that improves in Lecture 5 are measured on the same sequences.
-"""
-
-from __future__ import annotations
-
-from pathlib import Path
-
-import torch
-
-# %% 01. The DNA alphabet
-#
-# Chapters 4 and 5 run every worked example on DNA, so the vocabulary is four letters. The mask
-# token, when a lecture needs one, is appended as index 4 by that lecture, not here, because masked
-# diffusion has a mask and uniform diffusion does not.
-
-DNA = "ACGT"  # Token 0 is A, 1 is C, 2 is G, 3 is T.
-DNA_SIZE = len(DNA)  # K = 4, the alphabet size the chapters call |V|.
-MOTIF = "ACGT"  # The planted motif, used as the label in the guidance examples.
-
-
-def encode_dna(sequences: list[str]) -> torch.Tensor:
-    """Turn a list of strings over ACGT into a long tensor of shape [N, L]."""
-    index = {letter: i for i, letter in enumerate(DNA)}  # Letter to token id.
-    rows = [[index[letter] for letter in sequence] for sequence in sequences]
-    return torch.tensor(rows, dtype=torch.long)  # [N, L] token ids.
-
-
-def decode_dna(tokens: torch.Tensor) -> list[str]:
-    """Turn a long tensor of shape [N, L] back into strings over ACGT."""
-    return ["".join(DNA[int(i)] for i in row) for row in tokens]  # One string per row.
-
-
-# %% 02. Synthetic DNA with structure a small model can actually learn
-#
-# A uniform random sequence has nothing to learn, so the sampler would look the same before and
-# after training. These sequences carry two pieces of structure instead: a GC-rich background, and
-# the motif ACGT planted at a random position in a fraction of the sequences. The GC fraction then
-# gives a label for the conditional examples, and the motif gives a property to steer toward.
-
-
-def synthetic_dna(
-    num_sequences: int = 512,
-    length: int = 16,
-    motif_fraction: float = 0.75,
-    gc_bias: float = 0.6,
-    seed: int = 0,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Draw synthetic DNA with a planted motif.
-
-    Returns the tokens as a long tensor of shape [N, L] and a binary label of shape [N] that is
-    one when the sequence has GC fraction at least 0.6. The label is a deterministic function of
-    the sequence, not a measurement of anything biological.
-    """
-    generator = torch.Generator().manual_seed(seed)  # Local generator; global seed untouched.
-
-    # Background letters: C and G each get gc_bias/2, A and T share the rest.
-    background = torch.tensor([
-        (1.0 - gc_bias) / 2.0,  # A
-        gc_bias / 2.0,  # C
-        gc_bias / 2.0,  # G
-        (1.0 - gc_bias) / 2.0,  # T
-    ])
-    tokens = torch.multinomial(
-        background.expand(num_sequences * length, DNA_SIZE),  # One draw per position.
-        num_samples=1,
-        replacement=True,
-        generator=generator,
-    ).view(num_sequences, length)  # [N, L] background sequence.
-
-    # Plant the motif in a fraction of the sequences, at a position drawn uniformly.
-    motif = encode_dna([MOTIF])[0]  # [4] token ids of ACGT.
-    carries_motif = torch.rand(num_sequences, generator=generator) < motif_fraction  # [N] bool.
-    # The motif starts at one of three places, not anywhere, which leaves it a pattern a
-    # small model can learn from a few hundred sequences while still varying its position. A start
-    # is kept only when the whole motif fits, so a short sequence simply has fewer of them.
-    room = length - len(MOTIF)  # Last index at which the motif still fits.
-    choices = torch.tensor([s for s in (2, 6, 10) if s <= room] or [0])
-    starts = choices[torch.randint(0, len(choices), (num_sequences,), generator=generator)]
-    for row in torch.nonzero(carries_motif).flatten().tolist():
-        start = int(starts[row])
-        tokens[row, start:start + len(MOTIF)] = motif  # Overwrite those positions.
-
-    gc = ((tokens == 1) | (tokens == 2)).to(torch.get_default_dtype()).mean(dim=1)  # [N] fraction of C and G.
-    labels = (gc >= 0.6).long()  # [N] the toy label used by the conditional examples.
-    return tokens, labels
-
-
-def dna_splits(
-    num_sequences: int = 512, length: int = 16, seed: int = 0, validation: float = 0.2,
-) -> dict[str, torch.Tensor]:
-    """The same synthetic DNA, split into a training and a validation set."""
-    tokens, labels = synthetic_dna(num_sequences, length, seed=seed)
-    cut = int(num_sequences * (1.0 - validation))  # Index where validation begins.
-    return {
-        "train_tokens": tokens[:cut], "train_labels": labels[:cut],
-        "valid_tokens": tokens[cut:], "valid_labels": labels[cut:],
-    }
-
-
-def write_dna_tsv(path: str | Path, tokens: torch.Tensor, labels: torch.Tensor) -> None:
-    """Write sequences and labels as a two-column TSV, the format ``--data`` reads."""
-    lines = ["sequence\tlabel"]
-    lines += [f"{s}\t{int(y)}" for s, y in zip(decode_dna(tokens), labels)]
-    Path(path).write_text("\n".join(lines) + "\n")
-
-
-def read_dna_tsv(path: str | Path) -> tuple[torch.Tensor, torch.Tensor]:
-    """Read a two-column TSV of sequences and labels."""
-    rows = Path(path).read_text().strip().split("\n")[1:]  # Skip the header line.
-    sequences = [row.split("\t")[0] for row in rows]
-    labels = [int(row.split("\t")[1]) for row in rows]
-    return encode_dna(sequences), torch.tensor(labels, dtype=torch.long)
-
-
-# %% 03. Two-dimensional densities
-#
-# The continuous lectures need a target we can plot and whose exact density we know. Four Gaussians
-# at the corners of a square give both: the modes are far enough apart that a sampler that collapses
-# is obvious, and the log density is a closed form the reports can quote.
-
-GAUSSIAN_CENTERS = torch.tensor([[2.0, 2.0], [2.0, -2.0], [-2.0, 2.0], [-2.0, -2.0]])
-GAUSSIAN_STD = 0.35  # Standard deviation of each component.
-
-
-def gaussian_mixture_2d(num_points: int = 2048, seed: int = 0) -> torch.Tensor:
-    """Draw from four equally weighted Gaussians at the corners of a square. Returns [N, 2]."""
-    generator = torch.Generator().manual_seed(seed)
-    which = torch.randint(0, len(GAUSSIAN_CENTERS), (num_points,), generator=generator)  # [N].
-    noise = torch.randn(num_points, 2, generator=generator) * GAUSSIAN_STD  # [N, 2].
-    return GAUSSIAN_CENTERS[which] + noise  # [N, 2] points.
-
-
-def gaussian_mixture_log_density(points: torch.Tensor) -> torch.Tensor:
-    """Exact log density of the four-Gaussian mixture at each row of ``points`` ([N, 2] -> [N])."""
-    diff = points[:, None, :] - GAUSSIAN_CENTERS.to(points)  # [N, 4, 2] offset to each center.
-    squared = (diff ** 2).sum(dim=-1)  # [N, 4] squared distance to each center.
-    log_component = (
-        -0.5 * squared / GAUSSIAN_STD ** 2  # Gaussian exponent.
-        - torch.log(torch.tensor(2.0 * torch.pi * GAUSSIAN_STD ** 2))  # Normalizer in 2-D.
-    )  # [N, 4].
-    weight = torch.log(torch.tensor(1.0 / len(GAUSSIAN_CENTERS)))  # Equal mixture weights.
-    return torch.logsumexp(log_component + weight, dim=-1)  # [N] mixture log density.
-
-
-def two_moons(num_points: int = 2048, noise: float = 0.08, seed: int = 0) -> torch.Tensor:
-    """Two interleaving half circles, the standard test for a sampler that must bend. [N, 2]."""
-    generator = torch.Generator().manual_seed(seed)
-    half = num_points // 2
-    angle = torch.rand(half, generator=generator) * torch.pi  # [half] angle along each arc.
-    upper = torch.stack([torch.cos(angle), torch.sin(angle)], dim=1)  # Upper half circle.
-    lower = torch.stack([1.0 - torch.cos(angle), 0.5 - torch.sin(angle)], dim=1)  # Lower arc.
-    points = torch.cat([upper, lower], dim=0)  # [2*half, 2].
-    points = points + noise * torch.randn(points.shape, generator=generator)  # Blur the arcs.
-    return (points - points.mean(0)) / points.std(0)  # Standardize to roughly unit scale.
-
-
-def checkerboard(num_points: int = 2048, seed: int = 0) -> torch.Tensor:
-    """A checkerboard of eight squares, a density with disconnected support. [N, 2]."""
-    generator = torch.Generator().manual_seed(seed)
-    x = torch.rand(num_points, generator=generator) * 4.0 - 2.0  # [N] in [-2, 2].
-    y = torch.rand(num_points, generator=generator) - torch.randint(
-        0, 2, (num_points,), generator=generator,
-    ).to(torch.get_default_dtype()) * 2.0  # [N] in [-2, 1] before the shift below.
-    y = y + torch.floor(x) % 2.0  # Shift alternate columns so the squares interleave.
-    return torch.stack([x, y * 2.0], dim=1)  # [N, 2] points.
-
-
-# %% 04. Short phrases, for the discrete examples whose tokens are words
-#
-# Chapter 6 runs its categorical flow maps on four-token phrases. The grammar is small enough that
-# a reader can tell at a glance whether a generated phrase is in the language.
-
-_SUBJECTS = ["red", "blue", "green", "small"]
-_OBJECTS = ["circle", "square", "arrow", "dot"]
-_VERBS = ["moves", "turns", "drifts", "jumps"]
-_DIRECTIONS = ["left", "right", "up", "down"]
-
-
-def phrase_vocabulary() -> list[str]:
-    """The token vocabulary of the phrase language, in a fixed order."""
-    return _SUBJECTS + _OBJECTS + _VERBS + _DIRECTIONS  # 16 tokens, four per slot.
-
-
-def phrases(num_phrases: int = 512, seed: int = 0) -> tuple[torch.Tensor, list[str]]:
-    """Draw four-token phrases such as ``red circle moves left``. Returns [N, 4] and the vocabulary.
-
-    Each slot draws from its own quarter of the vocabulary, so a correct phrase is one whose
-    position i holds a token from group i. That makes the training support easy to check.
-    """
-    generator = torch.Generator().manual_seed(seed)
-    vocabulary = phrase_vocabulary()
-    groups = [_SUBJECTS, _OBJECTS, _VERBS, _DIRECTIONS]
-    columns = []
-    for slot, group in enumerate(groups):
-        offset = slot * len(group)  # Where this slot's tokens start in the vocabulary.
-        picks = torch.randint(0, len(group), (num_phrases,), generator=generator)
-        columns.append(picks + offset)  # [N] token ids inside the full vocabulary.
-    return torch.stack(columns, dim=1), vocabulary  # [N, 4] and the 16 token strings.
-
-
-# %% 05. MNIST and ESM-2, the two datasets that download
-#
-# Both are optional. A script reaches them only when the reader passes --dataset mnist or
-# --dataset esm2, and every such script also runs on the offline data above.
-
-
-def mnist_loader(batch_size: int = 128, root: str | Path = "data", train: bool = True):
-    """MNIST as a DataLoader of images normalized to [-1, 1]. Downloads on first use."""
-    from torch.utils.data import DataLoader  # Imported here so the module loads without torchvision.
-    from torchvision import datasets
-    from torchvision.transforms import v2
-
-    transform = v2.Compose([
-        v2.ToImage(),  # Channel-first image: [1, 28, 28].
-        v2.ToDtype(torch.float32, scale=True),  # Integer pixels 0-255 to floats in [0, 1].
-        v2.Normalize(mean=(0.5,), std=(0.5,)),  # Then [0, 1] to [-1, 1].
-    ])
-    dataset = datasets.MNIST(root=root, train=train, download=True, transform=transform)
-    return DataLoader(dataset, batch_size=batch_size, shuffle=train, num_workers=0)
-
-
-def toy_images(num_images: int = 256, size: int = 28, seed: int = 0) -> torch.Tensor:
-    """Images shaped like MNIST, drawn from a two-mode distribution of bars and blobs.
-
-    These stand in for MNIST wherever a test has to run without a download. They have the same
-    shape and the same [-1, 1] range, so the identical training and sampling code paths run.
-    """
-    generator = torch.Generator().manual_seed(seed)
-    images = torch.full((num_images, 1, size, size), -1.0)  # Background at the low end of [-1, 1].
-    grid = torch.arange(size)
-    for i in range(num_images):
-        if i % 2 == 0:  # A vertical bar.
-            column = int(torch.randint(6, size - 6, (1,), generator=generator))
-            images[i, 0, 4:size - 4, column - 2:column + 2] = 1.0
-        else:  # A round blob.
-            cy = int(torch.randint(8, size - 8, (1,), generator=generator))
-            cx = int(torch.randint(8, size - 8, (1,), generator=generator))
-            dy = (grid[:, None] - cy) ** 2
-            dx = (grid[None, :] - cx) ** 2
-            images[i, 0][(dy + dx) <= 25] = 1.0
-    images = images + 0.05 * torch.randn(images.shape, generator=generator)  # Light pixel noise.
-    return images.clamp(-1.0, 1.0)  # [N, 1, size, size] in the training range.
-
-
-PEPTIDES = [
-    "KWKLFKKIGAVLKVL", "GIGKFLHSAKKFGKA", "FLPIIAKLLSGLL", "ILPWKWPWWPWRR",
-    "KRWWKWWRR", "GLLSKLWKKVFKAFKKFLKK", "ACDEFGHIKLMNPQR", "RRWWRRWRR",
-    "LLGDFFRKSKEKIGKE", "KIAKVALKAL", "VRRFPWWWPFLRR", "GWLKKIGKKIERVGQ",
-]  # Short antimicrobial-peptide-like sequences, used when ESM-2 is unavailable.
-
-
-def esm2_embeddings(
-    sequences: list[str] | None = None,
-    model_name: str = "facebook/esm2_t33_650M_UR50D",
-    max_length: int = 32,
-    device: torch.device | str = "cpu",
-):
-    """Residue embeddings from a frozen ESM-2 encoder. Downloads the checkpoint on first use.
-
-    Returns the embeddings [N, L, D], the attention mask [N, L], the tokenizer and the masked
-    language-model head, which Lecture 3 uses to decode generated latents back to amino acids.
-    """
-    from transformers import AutoModel, AutoModelForMaskedLM, AutoTokenizer
-
-    sequences = sequences or PEPTIDES
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    encoder = AutoModel.from_pretrained(model_name).to(device).eval()
-    head = AutoModelForMaskedLM.from_pretrained(model_name).lm_head.to(device).eval()
-    for parameter in encoder.parameters():
-        parameter.requires_grad = False  # The encoder is frozen; only the denoiser trains.
-
-    batch = tokenizer(
-        sequences, padding="max_length", truncation=True,
-        max_length=max_length, return_tensors="pt",
-    ).to(device)
-    with torch.no_grad():
-        embeddings = encoder(**batch).last_hidden_state  # [N, L, D] residue embeddings.
-    return embeddings, batch["attention_mask"], tokenizer, head
-
-
-def toy_embeddings(
-    num_sequences: int = 128, length: int = 16, width: int = 32, seed: int = 0,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Embeddings shaped like ESM-2 output, for the offline path of the latent examples.
-
-    Each sequence is a smooth low-frequency curve along the residue axis plus noise, so neighbouring
-    positions are correlated the way real residue embeddings are. Returns [N, L, D] and a mask.
-    """
-    generator = torch.Generator().manual_seed(seed)
-    phase = torch.rand(num_sequences, 1, width, generator=generator) * 2.0 * torch.pi
-    frequency = torch.rand(num_sequences, 1, width, generator=generator) * 0.6 + 0.1
-    position = torch.arange(length).view(1, length, 1).to(torch.get_default_dtype())  # [1, L, 1].
-    signal = torch.sin(frequency * position + phase)  # [N, L, D] smooth along the residue axis.
-    noise = 0.2 * torch.randn(num_sequences, length, width, generator=generator)
-    mask = torch.ones(num_sequences, length)  # No padding in the toy set.
-    return signal + noise, mask
-
-
-def standardize(embeddings: torch.Tensor, mask: torch.Tensor):
-    """Center and scale each embedding coordinate over the valid residues.
-
-    Returns the standardized embeddings and the mean and standard deviation, which the decoder
-    needs in order to undo the transformation before reading amino acids.
-    """
-    valid = mask.bool().unsqueeze(-1)  # [N, L, 1] true at a real residue.
-    count = valid.sum()  # Number of valid residue-coordinate entries.
-    mean = (embeddings * valid).sum(dim=(0, 1)) / count  # [D] per-coordinate mean.
-    variance = ((embeddings - mean) ** 2 * valid).sum(dim=(0, 1)) / count  # [D].
-    std = variance.sqrt().clamp(min=1e-6)  # [D], clamped so a dead coordinate cannot divide by zero.
-    return (embeddings - mean) / std, mean, std

cis6270/metrics.py

@@ -1,100 +0,0 @@
-"""Measurements of a batch of generated samples.
-
-None of these is a benchmark. They are small, exactly computable quantities that tell a reader
-whether a sampler learned the structure planted in the data: the right letter frequencies, the
-planted motif, the right number of modes. A number here is a diagnostic, not a score.
-"""
-
-from __future__ import annotations
-
-import torch
-
-from cis6270.data import GAUSSIAN_CENTERS, MOTIF, decode_dna
-
-
-# %% 01. Measurements on DNA
-
-
-def gc_fraction(tokens: torch.Tensor) -> float:
-    """Fraction of positions holding C or G, averaged over the batch. The training data sit at 0.577."""
-    return float(((tokens == 1) | (tokens == 2)).to(torch.get_default_dtype()).mean())
-
-
-def motif_fraction(tokens: torch.Tensor, motif: str = MOTIF) -> float:
-    """Fraction of sequences that contain the planted motif anywhere."""
-    sequences = decode_dna(tokens)
-    return sum(motif in sequence for sequence in sequences) / max(len(sequences), 1)
-
-
-def unique_fraction(tokens: torch.Tensor) -> float:
-    """Fraction of generated sequences that are distinct, which catches a collapsed sampler.
-
-    The comparison is on the token ids themselves, so this works for any vocabulary, not only the
-    four DNA letters.
-    """
-    rows = [tuple(row.tolist()) for row in tokens]  # One tuple per sequence.
-    return len(set(rows)) / max(len(rows), 1)
-
-
-def token_entropy(tokens: torch.Tensor, vocab_size: int) -> float:
-    """Entropy in nats of the letter frequencies, at most log K. Four letters give log 4 = 1.386."""
-    counts = torch.bincount(tokens.flatten(), minlength=vocab_size).to(torch.get_default_dtype())
-    probabilities = counts / counts.sum().clamp(min=1)
-    nonzero = probabilities[probabilities > 0]
-    return float(-(nonzero * nonzero.log()).sum())
-
-
-def support_fraction(tokens: torch.Tensor, training: torch.Tensor) -> float:
-    """Fraction of generated sequences that also appear in the training set."""
-    seen = {tuple(row.tolist()) for row in training}
-    return sum(tuple(row.tolist()) in seen for row in tokens) / max(len(tokens), 1)
-
-
-# %% 02. Measurements on two-dimensional samples
-
-
-def mode_report(points: torch.Tensor, centers: torch.Tensor | None = None) -> dict:
-    """Assign each point to its nearest mode and report the distance and the coverage.
-
-    ``nearest_distance`` near the data's own standard deviation means the points land on the modes.
-    ``mode_entropy`` near log 4 = 1.386 means all four modes are visited about equally.
-    """
-    centers = GAUSSIAN_CENTERS.to(points) if centers is None else centers.to(points)
-    distance = torch.cdist(points, centers)  # [N, M] distance to every center.
-    nearest, which = distance.min(dim=1)  # [N] distance and index of the closest mode.
-    counts = torch.bincount(which, minlength=len(centers)).to(torch.get_default_dtype())
-    fractions = counts / counts.sum().clamp(min=1)
-    nonzero = fractions[fractions > 0]
-    return {
-        "nearest_distance": float(nearest.mean()),
-        "modes_found": int((counts > 0).sum()),
-        "mode_entropy": float(-(nonzero * nonzero.log()).sum()),
-        "mode_fractions": [round(float(f), 3) for f in fractions],
-    }
-
-
-def mean_log_density(points: torch.Tensor) -> float:
-    """Average exact log density of the four-Gaussian target at the generated points."""
-    from cis6270.data import gaussian_mixture_log_density
-
-    return float(gaussian_mixture_log_density(points).mean())
-
-
-# %% 03. Measurements on a learned distribution over a small state space
-#
-# Several chapters work on a state space small enough to enumerate, where the exact answer is
-# available. Comparing an empirical histogram with that exact answer is the sharpest check the
-# course has.
-
-
-def total_variation(p: torch.Tensor, q: torch.Tensor) -> float:
-    """Half the sum of absolute differences between two distributions over the same states."""
-    return float(0.5 * (p - q).abs().sum())
-
-
-def empirical_distribution(tokens: torch.Tensor, vocab_size: int, length: int) -> torch.Tensor:
-    """Histogram of whole sequences over the K**L states, for a state space small enough to list."""
-    weights = vocab_size ** torch.arange(length - 1, -1, -1, device=tokens.device)  # Place values.
-    indices = (tokens * weights).sum(dim=1)  # [N] index of each sequence in the full state space.
-    counts = torch.bincount(indices, minlength=vocab_size ** length).to(torch.get_default_dtype())
-    return counts / counts.sum().clamp(min=1)  # [K**L] empirical distribution.

cis6270/nets.py

@@ -1,231 +0,0 @@
-"""The small networks every lecture reuses.
-
-Four networks cover the whole course. None of them is the subject of a lecture: a lecture changes
-the loss and the sampler, and keeps the network fixed, so that any difference in the samples comes
-from the mathematics and not from the architecture.
-
-| Network            | Input                      | Output                | Used by |
-| ------------------ | -------------------------- | --------------------- | ------- |
-| `TimeEmbedding`    | a time or step index       | a vector of width D   | all     |
-| `SequenceDenoiser` | tokens [B, L] and a time   | logits [B, L, K]      | 4, 5, 6 |
-| `VelocityMLP`      | a point [B, D] and a time  | a vector [B, D]       | 2, 3, 6, 7 |
-| `TwoTimeMLP`       | a point, a start, an end   | a vector [B, D]       | 6 |
-| `SmallUNet`        | an image [B, 1, 28, 28]    | an image of the same  | 2, 3 |
-"""
-
-from __future__ import annotations
-
-import math
-
-import torch
-import torch.nn as nn
-
-
-# %% 01. Embedding a continuous time
-#
-# Every network in the course takes a time as well as a state. A raw scalar in [0, 1] is a poor
-# input to a linear layer, so it is first spread over many frequencies, exactly as in Section 3.28
-# of the notes, and then passed through a small MLP.
-
-
-def sinusoidal_embedding(t: torch.Tensor, width: int) -> torch.Tensor:
-    """Spread a time [B] over ``width`` sine and cosine frequencies. Returns [B, width]."""
-    half = width // 2  # Half the coordinates are sines and half cosines.
-    frequencies = torch.exp(
-        -math.log(10000.0) * torch.arange(half, device=t.device).to(torch.get_default_dtype()) / half
-    )  # [half] geometrically spaced frequencies.
-    angles = t.to(torch.get_default_dtype())[:, None] * frequencies[None, :]  # [B, half].
-    return torch.cat([torch.sin(angles), torch.cos(angles)], dim=1)  # [B, width].
-
-
-class TimeEmbedding(nn.Module):
-    """Sinusoidal features of the time, followed by a two-layer MLP."""
-
-    def __init__(self, width: int = 128, scale: float = 1000.0):
-        super().__init__()
-        self.width = width  # Width of the embedding the rest of the network receives.
-        self.scale = scale  # Times in [0, 1] are multiplied by this before the sines.
-        self.mlp = nn.Sequential(
-            nn.Linear(width, width), nn.SiLU(), nn.Linear(width, width),
-        )
-
-    def forward(self, t: torch.Tensor) -> torch.Tensor:
-        """t: [B] -> [B, width]."""
-        return self.mlp(sinusoidal_embedding(t * self.scale, self.width))  # [B, width].
-
-
-# %% 02. A denoiser over token sequences
-#
-# Lectures 4, 5 and 6 all train the same shape of network: it reads a corrupted sequence and a time,
-# and produces a distribution over the vocabulary at every position. The notes call it the denoiser
-# d_theta(. | z, t). Attention lets the prediction at one position depend on every other position,
-# which is what makes the factorized reverse step of Chapter 4 worth taking.
-
-
-class SequenceDenoiser(nn.Module):
-    """A small transformer that maps a token sequence and a time to per-position logits."""
-
-    def __init__(
-        self,
-        vocab_size: int,
-        length: int,
-        width: int = 128,
-        depth: int = 2,
-        heads: int = 4,
-        num_classes: int = 0,
-    ):
-        super().__init__()
-        self.vocab_size = vocab_size  # K, including a mask token when the lecture adds one.
-        self.length = length  # L, the number of positions.
-        self.width = width  # D, the model width.
-        self.token_embedding = nn.Embedding(vocab_size, width)  # One vector per token id.
-        self.position_embedding = nn.Embedding(length, width)  # One vector per position.
-        self.time_embedding = TimeEmbedding(width)  # Shared across positions.
-        # A label embedding, used only by the conditional examples. Index num_classes is the null
-        # label that classifier-free guidance drops to.
-        self.label_embedding = (
-            nn.Embedding(num_classes + 1, width) if num_classes else None
-        )
-        layer = nn.TransformerEncoderLayer(
-            d_model=width, nhead=heads, dim_feedforward=4 * width,
-            batch_first=True, dropout=0.0, norm_first=True,
-        )
-        self.transformer = nn.TransformerEncoder(
-            layer, num_layers=depth, enable_nested_tensor=False,
-        )
-        self.head = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, vocab_size))
-
-    def forward(
-        self, tokens: torch.Tensor, t: torch.Tensor, labels: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        """tokens: [B, L], t: [B], labels: [B] or None -> logits [B, L, K]."""
-        positions = torch.arange(tokens.shape[1], device=tokens.device)  # [L].
-        h = (
-            self.token_embedding(tokens)  # [B, L, D] what each position currently holds.
-            + self.position_embedding(positions)[None, :, :]  # [1, L, D] where it sits.
-            + self.time_embedding(t)[:, None, :]  # [B, 1, D] how far along the path we are.
-        )
-        if labels is not None and self.label_embedding is not None:
-            h = h + self.label_embedding(labels)[:, None, :]  # [B, 1, D] the condition.
-        h = self.transformer(h)  # [B, L, D] every position now sees every other.
-        return self.head(h)  # [B, L, K] logits over the vocabulary at each position.
-
-
-# %% 03. A velocity network on a vector state
-#
-# Lectures 2, 3, 6 and 7 learn a field on R^d: a velocity, a score, a noise prediction or a drift.
-# The same MLP serves for all of them, because the only difference is what the loss regresses it
-# onto.
-
-
-class VelocityMLP(nn.Module):
-    """An MLP from a point and a time to a vector of the same width as the point."""
-
-    def __init__(self, dim: int = 2, width: int = 128, depth: int = 3, num_classes: int = 0):
-        super().__init__()
-        self.dim = dim  # d, the width of the state.
-        self.time_embedding = TimeEmbedding(width, scale=1.0)  # Times here already live in [0, 1].
-        self.label_embedding = (
-            nn.Embedding(num_classes + 1, width) if num_classes else None
-        )
-        self.input = nn.Linear(dim, width)  # Lift the state to the hidden width.
-        layers: list[nn.Module] = []
-        for _ in range(depth):
-            layers += [nn.Linear(width, width), nn.SiLU()]
-        self.body = nn.Sequential(*layers)
-        self.output = nn.Linear(width, dim)  # Back down to one vector per coordinate.
-
-    def forward(
-        self, x: torch.Tensor, t: torch.Tensor, labels: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        """x: [B, d], t: [B], labels: [B] or None -> [B, d]."""
-        h = self.input(x) + self.time_embedding(t)  # [B, width] state and time together.
-        if labels is not None and self.label_embedding is not None:
-            h = h + self.label_embedding(labels)  # [B, width] the condition.
-        return self.output(self.body(h))  # [B, d].
-
-
-# %% 04. A two-time network, for the flow maps of Lecture 6
-#
-# A flow map carries a state from a start time s to an end time t, so its network takes both. The
-# identity axiom X_{s,s} = id is free if the map is written as the state plus the elapsed time times
-# a network output, which is what `flow_map` below does.
-
-
-class TwoTimeMLP(nn.Module):
-    """An MLP from a point, a start time and an end time to a vector."""
-
-    def __init__(self, dim: int = 2, width: int = 128, depth: int = 3):
-        super().__init__()
-        self.dim = dim
-        self.start_embedding = TimeEmbedding(width, scale=1.0)  # Embedding of s.
-        self.end_embedding = TimeEmbedding(width, scale=1.0)  # Embedding of t.
-        self.input = nn.Linear(dim, width)
-        layers: list[nn.Module] = []
-        for _ in range(depth):
-            layers += [nn.Linear(width, width), nn.SiLU()]
-        self.body = nn.Sequential(*layers)
-        self.output = nn.Linear(width, dim)
-
-    def forward(self, x: torch.Tensor, s: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
-        """The average velocity v_{s,t}(x). x: [B, d], s: [B], t: [B] -> [B, d]."""
-        h = self.input(x) + self.start_embedding(s) + self.end_embedding(t)  # [B, width].
-        return self.output(self.body(h))  # [B, d].
-
-    def flow_map(self, x: torch.Tensor, s: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
-        """X_{s,t}(x) = x + (t - s) v_{s,t}(x), so the identity axiom holds by construction."""
-        elapsed = (t - s)[:, None]  # [B, 1] elapsed time, broadcast over coordinates.
-        return x + elapsed * self.forward(x, s, t)  # [B, d].
-
-
-# %% 05. A small U-Net for 28-by-28 images
-#
-# Lectures 2 and 3 generate images. The network is deliberately small, about half a million
-# parameters, so that a few epochs on a laptop CPU already produce recognizable digits.
-
-
-def _conv_block(in_channels: int, out_channels: int) -> nn.Sequential:
-    """Two 3-by-3 convolutions with a SiLU after each, keeping the height and width."""
-    return nn.Sequential(
-        nn.Conv2d(in_channels, out_channels, 3, padding=1), nn.SiLU(),
-        nn.Conv2d(out_channels, out_channels, 3, padding=1), nn.SiLU(),
-    )
-
-
-class SmallUNet(nn.Module):
-    """A three-scale U-Net that maps an image and a time to an image of the same shape.
-
-    The time enters as an extra constant channel, which is the simplest conditioning that works and
-    is what the Lecture 2 script uses.
-    """
-
-    def __init__(self, base_channels: int = 32, in_channels: int = 1):
-        super().__init__()
-        c = base_channels  # Widths are c, 2c, 4c.
-        self.encoder1 = _conv_block(in_channels + 1, c)  # Image plus the time channel.
-        self.encoder2 = _conv_block(c, 2 * c)
-        self.middle = _conv_block(2 * c, 4 * c)
-        self.decoder2 = _conv_block(6 * c, 2 * c)  # Upsampled 4c joined with the 2c skip.
-        self.decoder1 = _conv_block(3 * c, c)  # Upsampled 2c joined with the c skip.
-        self.output = nn.Conv2d(c, in_channels, 1)  # One prediction per pixel.
-        self.pool = nn.MaxPool2d(2)  # Halve the height and width.
-
-    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
-        """x: [B, 1, 28, 28], t: [B] -> [B, 1, 28, 28]."""
-        import torch.nn.functional as F  # Local import keeps the module header short.
-
-        time_channel = t[:, None, None, None].expand(x.shape[0], 1, *x.shape[2:])  # [B, 1, H, W].
-        h = torch.cat([x, time_channel], dim=1)  # [B, 2, H, W].
-        skip1 = self.encoder1(h)  # [B, c, 28, 28].
-        skip2 = self.encoder2(self.pool(skip1))  # [B, 2c, 14, 14].
-        h = self.middle(self.pool(skip2))  # [B, 4c, 7, 7] the compressed representation.
-        h = F.interpolate(h, size=skip2.shape[-2:], mode="nearest")  # 7 to 14.
-        h = self.decoder2(torch.cat([h, skip2], dim=1))  # Rejoin the 14-by-14 detail.
-        h = F.interpolate(h, size=skip1.shape[-2:], mode="nearest")  # 14 to 28.
-        h = self.decoder1(torch.cat([h, skip1], dim=1))  # Rejoin the 28-by-28 detail.
-        return self.output(h)  # [B, 1, 28, 28].
-
-
-def count_parameters(model: nn.Module) -> int:
-    """Number of trainable weights and biases, printed by every script at startup."""
-    return sum(p.numel() for p in model.parameters() if p.requires_grad)

cis6270/runner.py

@@ -1,173 +0,0 @@
-"""Seeding, device selection, the training loop, checkpoints and the shared command line.
-
-Every lecture script ends in the same four lines: parse the arguments, build the model, call
-``train`` with a step function, and print a report. The step function is the only part that differs
-between frameworks, and it is the part the lecture is about.
-"""
-
-from __future__ import annotations
-
-import argparse
-import json
-import random
-import time
-from pathlib import Path
-from typing import Callable
-
-import torch
-
-
-# %% 01. Reproducibility and hardware
-
-
-def set_seed(seed: int) -> None:
-    """Seed Python, NumPy through torch, and torch itself, so a run repeats exactly."""
-    random.seed(seed)
-    torch.manual_seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(seed)
-
-
-def pick_device(requested: str = "auto") -> torch.device:
-    """Choose CUDA, then Apple MPS, then CPU, unless the reader names one."""
-    if requested != "auto":
-        return torch.device(requested)
-    if torch.cuda.is_available():
-        return torch.device("cuda")
-    if torch.backends.mps.is_available():
-        return torch.device("mps")
-    return torch.device("cpu")
-
-
-# %% 02. The command line every script shares
-#
-# A reader who has learned the flags once knows them for all fifty scripts. A script adds its own
-# flags on top of these with `parser.add_argument`.
-
-
-def common_parser(description: str) -> argparse.ArgumentParser:
-    """The flags shared by every script in the course."""
-    parser = argparse.ArgumentParser(
-        description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-    )
-    parser.add_argument("--steps", type=int, default=600, help="optimizer steps")
-    parser.add_argument("--batch-size", type=int, default=64, help="examples per step")
-    parser.add_argument("--lr", type=float, default=1e-3, help="learning rate")
-    parser.add_argument("--width", type=int, default=128, help="hidden width of the network")
-    parser.add_argument("--length", type=int, default=16, help="sequence length L")
-    parser.add_argument("--samples", type=int, default=16, help="sequences or points to generate")
-    parser.add_argument("--sample-steps", type=int, default=64, help="integration steps")
-    parser.add_argument("--seed", type=int, default=0, help="random seed")
-    parser.add_argument("--device", type=str, default="auto", help="auto, cpu, cuda or mps")
-    parser.add_argument("--out", type=str, default="", help="directory for the run's output")
-    parser.add_argument("--quiet", action="store_true", help="print only the final report")
-    return parser
-
-
-# %% 03. The training loop
-#
-# One loop for the whole course. The caller passes a function that takes the step index and returns
-# a scalar loss with a graph attached; everything else, the optimizer, the logging and the timing,
-# is the same in every lecture.
-
-
-def train(
-    model: torch.nn.Module,
-    loss_fn: Callable[[int], torch.Tensor],
-    steps: int = 600,
-    lr: float = 1e-3,
-    log_every: int = 100,
-    quiet: bool = False,
-    grad_clip: float = 1.0,
-    extra_parameters: list[torch.nn.Parameter] | None = None,
-) -> list[float]:
-    """Run ``steps`` optimizer steps on ``loss_fn`` and return the loss at every step.
-
-    ``loss_fn(step)`` draws its own batch, so a lecture that needs fresh noise at every step gets it
-    without a DataLoader. ``extra_parameters`` covers the rare case of a second trained object, such
-    as the bias force of Chapter 7.
-    """
-    parameters = list(model.parameters()) + list(extra_parameters or [])
-    optimizer = torch.optim.AdamW(parameters, lr=lr, weight_decay=0.0)
-    losses: list[float] = []
-    started = time.time()
-    model.train()
-    for step in range(steps):
-        loss = loss_fn(step)  # The lecture's own objective.
-        optimizer.zero_grad(set_to_none=True)
-        loss.backward()
-        torch.nn.utils.clip_grad_norm_(parameters, grad_clip)  # Keep a bad batch from derailing.
-        optimizer.step()
-        losses.append(float(loss.detach()))
-        if not quiet and log_every and (step + 1) % log_every == 0:
-            window = losses[-log_every:]  # Average over the last window, which is less jumpy.
-            print(f"  step {step + 1:5d}/{steps}  loss {sum(window) / len(window):.4f}")
-    if not quiet:
-        print(f"  trained {steps} steps in {time.time() - started:.1f}s")
-    model.eval()
-    return losses
-
-
-# %% 04. Saving a run
-#
-# A run writes three files: the weights, the configuration that produced them, and a short report of
-# what came out. No optimizer state, because none of these examples resumes training.
-
-
-def save_run(
-    out: str | Path,
-    model: torch.nn.Module | None = None,
-    config: dict | None = None,
-    report: dict | None = None,
-    losses: list[float] | None = None,
-    samples: list[str] | None = None,
-) -> Path | None:
-    """Write the weights, the configuration, the report, the loss curve and any samples."""
-    if not out:
-        return None
-    directory = Path(out)
-    directory.mkdir(parents=True, exist_ok=True)
-    if model is not None:
-        torch.save({"model": model.state_dict(), "config": config or {}},
-                   directory / "checkpoint.pt")
-    if config is not None:
-        (directory / "config.json").write_text(json.dumps(config, indent=2) + "\n")
-    if report is not None:
-        (directory / "report.json").write_text(json.dumps(report, indent=2) + "\n")
-    if losses is not None:
-        # Thin the curve to at most 200 points so the file stays small and readable.
-        stride = max(1, len(losses) // 200)
-        thinned = [round(v, 5) for v in losses[::stride]]
-        (directory / "losses.json").write_text(json.dumps(thinned) + "\n")
-    if samples is not None:
-        (directory / "samples.txt").write_text("\n".join(samples) + "\n")
-    return directory
-
-
-def load_run(out: str | Path, model: torch.nn.Module) -> dict:
-    """Read the weights back into ``model`` and return the configuration that trained them."""
-    payload = torch.load(Path(out) / "checkpoint.pt", map_location="cpu", weights_only=False)
-    model.load_state_dict(payload["model"])
-    model.eval()
-    return payload.get("config", {})
-
-
-# %% 05. Printing a report
-#
-# Every script ends with the same two-column block, so a reader can compare two methods by reading
-# two reports side by side.
-
-
-def print_report(title: str, report: dict, samples: list[str] | None = None) -> None:
-    """Print a titled block of name and value pairs, and a few samples."""
-    print(f"\n{title}")
-    print("-" * len(title))
-    for key, value in report.items():
-        if isinstance(value, float):
-            print(f"  {key:<34s} {value:.4f}")
-        else:
-            print(f"  {key:<34s} {value}")
-    if samples:
-        print("\n  samples")
-        for sample in samples[:8]:
-            print(f"    {sample}")