NumPy and vectorized thinking

Arrays: one block of numbers, many views

A NumPy ndarray is a contiguous block of memory plus metadata: a shape (a tuple of axis lengths), a dtype (every element has the same type, e.g. float32, int64, bool) and strides (how many bytes to step along each axis). Because the metadata is separate from the data, operations like transposing, reshaping and slicing create views — new metadata over the same memory — in constant time. That is a superpower (reshaping a million numbers is free) and a trap (modifying a slice modifies the original; call .copy() when you need independence).

Read shapes as meanings. A dataset of 60,000 grayscale digits is (60000, 28, 28): samples × rows × columns. A batch of 32 token embeddings of width 768 is (32, 768). A batch of 8 sequences of 128 tokens, each a 512-dimensional vector, is (8, 128, 512) — batch × time × features, the shape convention transformers use (7.2 Scaled dot-product and multi-head attention, fully derived). Half of all ML bugs are shape bugs, so print shapes relentlessly.

import numpy as np
X = np.random.default_rng(0).normal(size=(60000, 28, 28)).astype(np.float32)
X.shape, X.dtype, X.nbytes / 1e6        # (60000, 28, 28), float32, 188.16 MB
flat = X.reshape(60000, -1)             # (60000, 784): -1 means "infer this axis"
flat.base is X                          # True: a view, no copy

dtype choices are cost choices: float64 is NumPy's default and 8 bytes per number; deep learning runs in float32 (4 bytes) or smaller (4.12 Compute, hardware and scale basics). Integers for labels and token ids, booleans for masks.

Indexing and masks

Basic slicing works per axis: X[0] is the first image (28, 28), X[:, 10, :] is row 10 of every image (60000, 28), X[..., ::2] takes every other column. Boolean masking selects by condition: y[y == 7], X[y == 7].mean(axis=0) gives the average "7". Fancy indexing with an integer array gathers rows: E[ids] where E is a (V, d) embedding table and ids is (B, T) returns (B, T, d) — that single line is an embedding layer (6.3 Word embeddings). Fancy indexing copies; slicing views.

Broadcasting: the rule that replaces loops

Broadcasting lets arrays of different shapes combine element-wise. The rule: align the shapes from the right; two axes are compatible if they are equal or one of them is 1; a missing axis counts as 1. The result has the maximum length along each axis, and the length-1 axes are conceptually repeated without copying memory.

The last row is the classic mistake: standardizing columns is (X - X.mean(axis=0)) / X.std(axis=0) because the mean over axis 0 has shape (d,) which aligns with the last axis; standardizing rows needs keepdims=True so the mean is (n, 1). Use None (or np.newaxis) to insert axes: a[:, None] - b[None, :] builds every pairwise difference between vectors a and b in one expression.

Reductions along axes

sum, mean, max, argmax, std all take axis: the axis you name is the one that disappears. For logits of shape (B, C), logits.argmax(axis=1) is the predicted class per example, (B,). keepdims=True keeps the reduced axis with length 1 so the result broadcasts back — which is exactly what a numerically stable softmax needs:

Subtracting the maximum does not change the result (multiply numerator and denominator by $e^{-m}$) but prevents $e^{z_i}$ from overflowing to infinity when $z_i > 709$ in float64 or $> 88$ in float32.

def softmax(z, axis=-1):
    z = z - z.max(axis=axis, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=axis, keepdims=True)

Matrix products: @ and einsum

A @ B is matrix multiplication (np.matmul); A * B is element-wise. For 2-D arrays (n, k) @ (k, m) -> (n, m); for stacks, @ broadcasts over leading batch axes, so (8, 128, 64) @ (8, 64, 128) -> (8, 128, 128) computes eight attention-score matrices at once (7.2 Scaled dot-product and multi-head attention, fully derived). A linear layer with inputs $\mathbf{X}\in\mathbb{R}^{n\times d}$ and weights $\mathbf{W}\in\mathbb{R}^{k\times d}$ (PyTorch's convention, out × in) is

costing $2ndk$ floating-point operations. np.einsum spells out the index pattern: np.einsum("nd,kd->nk", X, W) says "sum over the shared index d". It shines when the operation has no name: einsum("bqd,bkd->bqk", Q, K) is batched $\mathbf{Q}\mathbf{K}^\top$; einsum("ii->", A) is a trace; einsum("bi,bi->b", a, b) is a per-row dot product.

Worked example: pairwise distances without loops

kNN (3.4 k-Nearest Neighbors) and k-means (3.11 Clustering: k-means, GMMs/EM, DBSCAN and hierarchical) need the distance from every query to every training point. For $\mathbf{a}_i, \mathbf{b}_j \in \mathbb{R}^d$:

The three terms have shapes (n, 1), (1, m) and (n, m), so broadcasting assembles the full $n\times m$ matrix from one matrix product:

def pairwise_sq_dists(A, B):               # A: (n, d), B: (m, d)
    a2 = (A ** 2).sum(axis=1, keepdims=True)   # (n, 1)
    b2 = (B ** 2).sum(axis=1)[None, :]         # (1, m)
    return np.maximum(a2 + b2 - 2 * A @ B.T, 0)  # clip tiny negatives from rounding

rng = np.random.default_rng(0)
A, B = rng.normal(size=(1000, 64)), rng.normal(size=(5000, 64))
D = pairwise_sq_dists(A, B)                # (1000, 5000) in ~10 ms
nearest = D.argmin(axis=1)                 # index of nearest B for each A

A double Python loop over 5 million pairs takes seconds to minutes; this takes milliseconds. The speed-up comes from three things: the loop runs in compiled C, the data is contiguous in memory so the CPU streams it, and the matrix product uses a tuned BLAS library that exploits SIMD instructions and cache blocking. The rule of thumb: if you wrote for i in range(n) over data, there is almost certainly an array expression that replaces it.

Randomness done right

Use the modern generator API: rng = np.random.default_rng(seed), then rng.normal(size=...), rng.integers(0, V, size=...), rng.permutation(n) for shuffling, rng.choice(n, size=k, replace=False) for sampling. Pass the generator around rather than seeding globally; it makes experiments reproducible (2.5 Tooling: Git, GPUs, experiments and the Hugging Face Hub).

Where this shows up later

• Shapes and broadcasting are identical in PyTorch — 2.4 PyTorch I: tensors, autograd and the training loop; batch × time × features is the transformer convention — 7.2 Scaled dot-product and multi-head attention, fully derived.
• Stable softmax reappears in cross-entropy with logits — 4.3 Loss functions catalog — and in attention — 7.2 Scaled dot-product and multi-head attention, fully derived.
• Pairwise distances power kNN and k-means — 3.4 k-Nearest Neighbors, 3.11 Clustering: k-means, GMMs/EM, DBSCAN and hierarchical; fancy indexing is the embedding lookup — 6.3 Word embeddings.
• The linear-layer formula and its $2ndk$ cost are the unit of FLOP accounting — 4.12 Compute, hardware and scale basics, 7.9 Scaling laws and compute-optimal training.