Python for ML, fast

Why Python, and how much of it

Machine-learning code is mostly orchestration: load data, reshape it, call a library that runs a compiled kernel on thousands of numbers at once, log a metric, repeat. Python is slow at tight loops but superb at orchestration, and its ecosystem (NumPy, pandas, PyTorch, scikit-learn, Hugging Face) is where the field lives. The consequence is a specific learning strategy: master the constructs that ML code uses constantly and ignore the rest until you need it. This lesson is that subset. If a construct below is new, type the snippet yourself and change one thing — reading alone does not build fluency.

Containers and comprehensions

Four containers cover almost everything. A list is an ordered, mutable sequence; a tuple is an immutable one (use it for fixed-size records like (height, width) and as dictionary keys); a dict maps keys to values in insertion order; a set holds unique items with $O(1)$ membership tests, which is why vocabularies and "seen" bookkeeping use sets.

tokens = ["the", "cat", "sat", "on", "the", "mat"]
vocab = sorted(set(tokens))                    # unique words, sorted
stoi = {w: i for i, w in enumerate(vocab)}     # dict comprehension: word -> index
ids = [stoi[w] for w in tokens]                # list comprehension: encode
pairs = list(zip(ids[:-1], ids[1:]))           # bigrams as tuples
counts = {}
for a, b in pairs:                             # tuple unpacking in the loop header
    counts[(a, b)] = counts.get((a, b), 0) + 1
top = sorted(counts.items(), key=lambda kv: -kv[1])[:3]

That block is a complete bigram language model's counting step (6.2 Classical NLP and n-gram language models, 6.8 Project: character-level language model). Notice the idioms: comprehensions instead of append-loops, enumerate and zip instead of index arithmetic, sorted(..., key=...) with a lambda, and dict.get(k, default). Slicing (x[:-1], x[::2]) returns copies for lists but views for NumPy arrays — a distinction that will matter in 2.2 NumPy and vectorized thinking.

Functions, closures and generators

Functions are first-class values. You will pass them as arguments (a loss function into a trainer), return them from other functions (a learning-rate schedule), and store them in dicts (a registry of models). Default arguments are evaluated once at definition time, so never use a mutable default like def f(x, acc=[]); use None and create the list inside. Keyword-only arguments (after a bare *) make calls self-documenting: train(model, data, *, lr=1e-3, epochs=10). *args and **kwargs forward arbitrary arguments, which is how wrappers and framework callbacks are written.

A generator produces values lazily with yield, keeping one item in memory at a time. Data pipelines are generators: a function that opens a 50 GB file and yields one example per line can be consumed by a training loop without ever loading the file. Generator expressions (f(x) for x in xs) are the lazy cousin of list comprehensions.

def batches(xs, size):
    for i in range(0, len(xs), size):
        yield xs[i:i + size]

def make_schedule(base_lr, warmup):
    def lr_at(step):                 # closure: remembers base_lr and warmup
        return base_lr * min(1.0, (step + 1) / warmup)
    return lr_at

lr = make_schedule(3e-4, warmup=100)
print(lr(0), lr(50), lr(500))        # 3e-06 1.53e-04 3e-04

Classes, dataclasses and typing

You will write few classes but read many. Two matter most. A plain class with init and methods is how PyTorch models are defined (nn.Module subclasses, 2.4 PyTorch I: tensors, autograd and the training loop). A @dataclass is the idiomatic way to hold configuration: it writes init, repr and equality for you, supports defaults, and is trivially serialized to JSON for experiment logs.

from dataclasses import dataclass, asdict

@dataclass
class TrainConfig:
    lr: float = 3e-4
    batch_size: int = 64
    epochs: int = 10
    model_name: str = "mlp"

cfg = TrainConfig(lr=1e-3)
print(asdict(cfg))        # {'lr': 0.001, 'batch_size': 64, 'epochs': 10, 'model_name': 'mlp'}

Type hints (float, list[int], dict[str, float], Callable, Optional[X] or X | None) are not enforced at runtime, but they are documentation your editor understands and they make library signatures readable: def forward(self, x: Tensor) -> Tensor tells you the shape story is in the docstring, not the types. Dunder methods you will meet: len and getitem (how a PyTorch Dataset works), call (why you call a model like a function), repr, and iter.

The small things that prevent big bugs

• f-strings for all formatting: f"epoch {epoch:03d} loss {loss:.4f}". The format specs (:.4f, :,, :>8) make logs readable.
• pathlib instead of string paths: Path("data") / "train.csv", .exists(), .glob("*.png"), .read_text(). It works identically on Windows and Linux.
• Context managers guarantee cleanup: with open(p) as f:, with torch.no_grad():. Write your own with @contextmanager when you need a timer or a temporary setting.
• Exceptions: catch the narrowest exception you can (except FileNotFoundError), never bare except:, and re-raise with context when you add information. A traceback reads bottom-up: the last line is the error, the frames above it are where it came from.
• Debugging: put breakpoint() on the line before the failure, run the script, and inspect variables (p var, pp var.shape, n for next, c to continue). In notebooks, %debug after an exception drops you into the failing frame. Printing shapes is the single most effective ML debugging habit.

Environments: uv, pip and conda

Every project gets its own environment so that one project's PyTorch version cannot break another's. In 2026 the fastest, simplest tool is uv: uv init myproj creates a pyproject.toml, uv add numpy pandas torch adds dependencies and locks exact versions in uv.lock, and uv run train.py executes inside the environment without activating anything. The lock file is what makes a project reproducible on another machine six months later. Use conda (or its faster drop-in, mamba) when you need non-Python binaries such as a specific CUDA toolkit, or when a research repo ships an environment.yml. Plain pip plus python -m venv .venv still works everywhere and is what you will see in older tutorials.

Notebooks versus scripts: notebooks (Jupyter, Colab) are for exploration — you want to look at data and iterate on a plot. Scripts are for anything you will run more than twice — training runs, data processing — because they are versionable, testable and restartable. The common failure mode is a notebook whose cells were run out of order; if a result surprises you, restart the kernel and run all.

Where this shows up later

• Comprehensions, dicts and generators build tokenizers and data pipelines — 6.1 Text to numbers: tokenization, 4.9 PyTorch II: real training pipelines.
• Classes with init/forward and getitem are PyTorch models and datasets — 2.4 PyTorch I: tensors, autograd and the training loop, 4.9 PyTorch II: real training pipelines.
• Dataclass configs and f-string logging are how experiments stay comparable — 2.5 Tooling: Git, GPUs, experiments and the Hugging Face Hub, 4.8 Training dynamics and debugging.
• Closures become learning-rate schedules and callbacks — 4.5 Optimizers and learning-rate schedules; context managers wrap no_grad evaluation — 2.4 PyTorch I: tensors, autograd and the training loop.