Vectorization and profiling

Every experiment you run per day is a function of how fast your code is, and research velocity compounds like interest. The gap between idiomatic and naive numerical Python is not 20% — it is two orders of magnitude, and it comes from exactly two mechanisms you can learn in an afternoon: the interpreter tax and the memory hierarchy. This lesson explains both, states the broadcasting rules precisely (once, so the guessing stops), and then teaches the only reliable optimization method ever discovered: measure first.

The 100×: where it actually comes from

import numpy as np
a, b = np.random.rand(10_000_000), np.random.rand(10_000_000)

def loop(a, b):
    out = np.empty_like(a)
    for i in range(len(a)):
        out[i] = a[i] * b[i]
    return out

# %timeit loop(a, b)     # ~2.5 s
# %timeit a * b          # ~15 ms — the same arithmetic, ~150× faster

Two taxes are being refunded. The interpreter tax: each trip through a Python loop body costs bytecode dispatch, type checks, and boxing floats into heap objects — tens of nanoseconds of overhead wrapping a sub-nanosecond multiply. a * b pays that overhead once and runs a compiled C loop over raw doubles, which the CPU further accelerates with SIMD (8+ multiplies per instruction). The memory tax: performance for big arrays is bounded by bytes moved, not FLOPs (the GPU lesson’s memory-bound regime — it rules CPUs too). RAM arrives in 64-byte cache lines, so touching memory sequentially gets 8 doubles per line and lets the prefetcher stream; striding or hopping wastes 7/8 of every fetch.

That second tax has a famous fingerprint: layout. NumPy arrays are row-major (C order) — rows contiguous. Summing a large matrix along rows vs columns is the same arithmetic with different stride patterns, and the timing difference (often 2–5× at sizes beyond cache) is pure memory hierarchy. The general habits: iterate/reduce along the contiguous axis when you can, know that transpose/swapaxes are free (they change strides, not data) until something forces a copy, and reach for arr.copy(order="C") when a hot loop will re-read a badly-laid-out view many times.

Broadcasting, stated precisely

The rules — all of them, there are only two: align shapes from the right; two dimensions are compatible if equal or either is 1, and size-1 dimensions are (virtually — no copy) stretched to match.

X = np.random.rand(500, 3)         # (500, 3)
mu = X.mean(axis=0)                # (3,)
Xc = X - mu                        # (500,3)−(3,) → aligns right → OK

col = np.random.rand(500)          # (500,)
# X - col                          # (500,3)−(500,): 3 vs 500 → ERROR (a good one)
Xr = X - col[:, None]              # (500,1) → stretches to (500,3): what you meant

D = ((A[:, None, :] - B[None, :, :]) ** 2).sum(-1)   # all pairwise distances,
# (n,1,d)−(1,m,d)→(n,m,d): the idiom behind every kernel/distance matrix

The habit that catches 90% of silent bugs: comment expected shapes and assert the ones that matter. The worst broadcasting outcome is not the error — errors are gifts — it is the (500,1) vs (500,) subtraction that succeeds with shape (500,500) and a wrong answer. (keepdims=True on reductions exists to prevent exactly this.) When an expression needs three Nones and a transpose, np.einsum("ij,kj->ik", A, B) says the same thing legibly — and everything in this section transfers verbatim to torch/JAX, where it becomes GPU speed as well as CPU speed.

Profiling: the discipline of not guessing

Decades of evidence, one sentence: programmers’ guesses about where time goes are usually wrong, and optimization without measurement is superstition. The working ladder:

• %timeit for micro-questions (is A faster than B?) — it handles warmup and repetition; never trust a single time.time() delta.
• cProfile → snakeviz for the macro map: which functions own the time. First finding, in most ML code: it’s the dataloader / IO / a pandas apply, not the math you were beautifying.
• line_profiler (@profile + kernprof) for the micro map: which line in the hot function. This is where the accidental O(n²) — a list.index in a loop, a repeated pd.concat, an unintended copy per iteration — gets caught red-handed.
• GPU work needs its own stopwatch: kernels launch asynchronously, so wall-clocking a line measures the launch, not the work — call torch.cuda.synchronize() around timings or use torch.profiler/Nsight, or your profile is fiction.

And when a loop is genuinely irreducible — dynamic programming, per-element state, simulation steps — the escape hatch is compilation, not cleverness: numba @njit for numeric loops (near-C speed for the cost of a decorator), Cython for library-grade code, or restructuring the algorithm so the loop moves into NumPy/torch after all. The priority order stands: correct → measured → vectorized → compiled, and only ever the measured-hot part. Readability is a feature; a 3× speedup of 2% of runtime is a negative-value patch.

Exercises