NumPy & Math Foundations

Every model in this roadmap is a long chain of operations on tables of numbers. Before backpropagation, attention or sampling, you need to be comfortable with four things, and this page explains each one with small examples:

• what a NumPy array is, and what its shape and axes mean;
• how broadcasting combines arrays of different shapes;
• why exp of a large number breaks, and the two tricks that fix it;
• two indexing tools that come back in every later section.

Arrays, shapes and axes

A NumPy array is a table of numbers, all of the same type. A list of numbers is a 1-D array, a table with rows and columns is a 2-D array, and you can keep going: a stack of tables is 3-D, and so on.

The shape says how many numbers there are along each direction. Each direction is called an axis, and axes are numbered from 0.

Example. x = np.array([[1, 2, 3], [4, 5, 6]]) is a table with 2 rows and 3 columns, so x.shape is (2, 3).

• Axis 0 goes down the rows (length 2). x.sum(axis=0) adds up each column, moving down: [1+4, 2+5, 3+6] = [5, 7, 9].
• Axis 1 goes across the columns (length 3). x.sum(axis=1) adds up each row, moving across: [1+2+3, 4+5+6] = [6, 15].

Summing along an axis makes that axis disappear: shape (2, 3) becomes (3,) or (2,).

Behind the scenes, an array is one block of memory plus two pieces of bookkeeping: its shape, and its strides, how many bytes to jump in memory to reach the next number along each axis. For the x above, stored as 8-byte numbers, the strides are (24, 8): the next column is 8 bytes away (one number), the next row 24 bytes away (three numbers). This is why operations such as transposing are instant: NumPy only changes the bookkeeping, not the numbers.

Think in arrays, not loops

An expression like x @ W + b (a matrix product, then an addition) runs as compiled loops inside NumPy, so it is often orders of magnitude (factors of 10, 100 or more) faster than the same arithmetic written as a Python for loop.

The rule for every exercise here: no Python loops over elements. If you catch yourself writing one, look for the array expression that does the same job.

Example. To add 1 to every number of x, do not loop over rows and columns: write x + 1. The result is [[2, 3, 4], [5, 6, 7]].

Always know the shape of every array, and write it next to the code, for example # (N, D) for $N$ rows of $D$ numbers. Most bugs in ML (machine learning) code are shape bugs, and NumPy often hides them instead of raising an error, because of broadcasting.

Broadcasting

Broadcasting is what NumPy does when two arrays of different shapes meet in an operation that works number by number (+, -, *, /). It follows three steps:

1. Write the two shapes one above the other, aligned to the right. If one shape is shorter, pad it on the left with 1s.
2. Compare them axis by axis. Two sizes fit together when they are equal, or when one of them is 1. Otherwise NumPy raises an error.
3. An axis of size 1 is repeated to match the other size, without actually copying any memory.

Example 1: a row added to every row. x has shape (2, 3) and b = [10, 20, 30] has shape (3,).

1. Aligned to the right: (2, 3) and (3,), so b is treated as shape (1, 3).
2. Axis 1: 3 and 3 are equal. Axis 0: 2 and 1, fine because one is 1.
3. b is repeated down the 2 rows. x + b is [[11, 22, 33], [14, 25, 36]].

Example 2: every pair. A column c = [[1], [2], [3]], shape (3, 1), plus a row r = [[10, 20]], shape (1, 2). Both axes have a 1, so the column is repeated across 2 columns and the row down 3 rows. c + r is the full $3 \times 2$ table of sums, [[11, 21], [12, 22], [13, 23]]: entry $(i, j)$ is c[i] + r[j].

Example 2 is the trick behind every "all pairs" computation. To turn a row of numbers into a column, insert an axis of size 1 with x[:, None] (None is the same as np.newaxis).

Example 3: the classic bug. You want to divide every row of x by its own sum, so that each row adds up to 1. The row sums x.sum(axis=1) are [6, 15], with shape (2,).

• x / x.sum(axis=1) aligns (2, 3) with (2,) to the right: 3 against 2, neither is 1, so NumPy raises operands could not be broadcast together. Worse, with a square table it would not raise: it would divide column $j$ by the sum of row $j$, and return wrong numbers without any warning.
• x / x.sum(axis=1, keepdims=True) works. With keepdims=True the summed axis is kept with size 1, so the sums have shape (2, 1): a column, repeated across. The result is [[1/6, 2/6, 3/6], [4/15, 5/15, 6/15]], about [[0.167, 0.333, 0.5], [0.267, 0.333, 0.4]].

This is why reductions such as sum and max take keepdims=True: their result keeps a size-1 axis, so it lines up with the array it came from.

Floating point and numerical stability

Computers store real numbers as floating-point numbers. NumPy's default, float64, keeps about 16 significant decimal digits, and its largest value is about $1.8 \times 10^{308}$. Anything bigger becomes inf (infinity), and anything too close to 0 becomes exactly 0:

• $e^{709} \approx 8.2 \times 10^{307}$ still fits, but $e^{710}$ is inf;
• $e^{-745} \approx 5 \times 10^{-324}$ is the smallest number left, and $e^{-746}$ is exactly 0.

This matters because models produce logits: raw scores, any real numbers, one per possible answer. The next step is almost always $e^{\text{logit}}$, and logits of a few hundred or a few thousand do happen. When they overflow, the training loss becomes nan ("not a number") and training stops making sense. Two tricks fix almost everything.

Trick 1: shift before softmax

Softmax turns a list of scores into probabilities: positive numbers that add up to 1. It takes $e$ to the power of each score, then divides by the total:

where:

• $x = (x_1, \ldots, x_n)$ is the list of scores and $x_i$ is one of them;
• $e^{x_i}$ is always positive, and bigger scores give bigger values;
• $\sum_j e^{x_j}$ is the total of all of them, so dividing by it makes the results add up to 1.

Example. For $x = (1, 2, 3)$: $e^1 \approx 2.718$, $e^2 \approx 7.389$, $e^3 \approx 20.086$, total $\approx 30.193$. Dividing each by the total gives $(0.090, 0.245, 0.665)$, which adds up to 1.

With $x = (1000, 1001)$ the same steps compute $e^{1000}$, which is inf, and inf / inf is nan. The fix: subtracting the same number $c$ from every score does not change the result,

because $e^{x_i - c} = e^{x_i} \cdot e^{-c}$, and the same factor $e^{-c}$ appears above and below the fraction, so it cancels. Choose $c$ = the largest score: then the largest exponent is $e^0 = 1$ and nothing can overflow.

Example. For $x = (1000, 1001)$, subtract $c = 1001$: the scores become $(-1, 0)$. Then $e^{-1} \approx 0.368$ and $e^0 = 1$, total $1.368$, and the softmax is $(0.269, 0.731)$, the correct answer.

Trick 2: log-sum-exp

Training works with the logarithm of probabilities, and needs the quantity "log of the sum of the exponentials", $\log \sum_i e^{x_i}$, called log-sum-exp. It overflows for the same reason, and it has its own fix: pull the largest score $m$ out of the sum.

where:

• $x_i$ are the scores;
• $m = \max_i x_i$ is the largest of them.

Why it works: after subtracting $m$, the largest exponent is $e^0 = 1$, so nothing overflows. And the sum contains that 1, so it is at least 1, and its log is never $\log 0$. This is why frameworks compute log_softmax(x) = x - logsumexp(x) directly, instead of log(softmax(x)).

Example. For $x = (1000, 1000)$: $m = 1000$, the shifted scores are $(0, 0)$, the sum is $e^0 + e^0 = 2$, so the result is $1000 + \log 2 \approx 1000.693$. The naive formula would compute $\log(\infty + \infty) = \infty$.

Two indexing tools

Picking one entry per row

Integer-array indexing gives NumPy two lists, one of row numbers and one of column numbers, and it pairs them up position by position. P[np.arange(N), targets] reads, from each row i, the column targets[i]. This is how you read the probability of the correct token out of a table of probabilities.

Example. P = [[0.2, 0.7, 0.1], [0.5, 0.1, 0.4]] and targets = [1, 2]. np.arange(2) is [0, 1], so P[[0, 1], [1, 2]] reads P[0, 1] and P[1, 2]: the result is [0.7, 0.4].

np.einsum: products and sums by letters

np.einsum lets you describe a computation by giving each axis a letter. Letters in the inputs that are missing from the output are summed over; the others are kept.

Example. The matrix product A @ B is np.einsum("ij,jk->ik", A, B): for each row i of A and column k of B, multiply along the shared letter j and add up. With A = [[1, 2], [3, 4]] and B = [[5, 6], [7, 8]], entry $(0, 0)$ is $1 \cdot 5 + 2 \cdot 7 = 19$, and the full result is [[19, 22], [43, 50]].

The same notation scales to the 4-D arrays of attention: np.einsum("bhid,bhjd->bhij", q, k) computes all the attention scores for every batch b and head h in one line. The letter d is missing from the output, so it is summed over.

References

• C. R. Harris et al., Array programming with NumPy, Nature 585, 357–362, 2020. How arrays, strides and broadcasting are designed.
• NumPy documentation, Broadcasting and numpy.einsum.
• D. Goldberg, What every computer scientist should know about floating-point arithmetic, ACM Computing Surveys 23(1), 5–48, 1991.
• P. Blanchard, D. J. Higham, N. J. Higham, Accurately computing the log-sum-exp and softmax functions, IMA Journal of Numerical Analysis 41(4), 2311–2330, 2021.
• I. Goodfellow, Y. Bengio, A. Courville, Deep Learning, MIT Press, 2016. Chapter 4: Numerical Computation.