Python and NumPy for math: arrays, broadcasting, plots and random numbers

Before this lesson: lesson 02 (indices, Σ-notation, $\R^n$ and $\R^{m \times n}$). You need no Python installation: the notebook runs in Colab or in your browser. Time: about 70 minutes, plus the notebook. This is a long lesson; you can stop after section 3.6 and come back.

---

Why this matters #

Every notebook on this site uses NumPy. The deep-learning libraries PyTorch and JAX copy its array model: a PyTorch tensor has a .shape, is indexed and sliced in the same way, and follows the same broadcasting rule. scikit-learn expects a data matrix X of shape (m, n): $m$ samples (rows) and $n$ features (columns). A batch of colour images has shape (B, C, H, W) (batch, channels, height, width), and "normalize each channel" is one broadcasting line.

The same ideas appear in every model: a linear layer is X @ W + b (a matrix product plus a broadcast bias), standardizing data is (X - X.mean(axis=0)) / X.std(axis=0), and the predicted class is logits.argmax(axis=1). If you can read these lines and say the shape of each part, you can read most ML code.

1. Predict first #

2. The picture #

If Python itself is new for you, read section 3.1 first and then come back.

An array is a grid of numbers with a shape. The notation of lesson 02 and NumPy are two views of the same objects:

Math	NumPy	Note
$x \in \R^n$	x.shape == (n,)	a 1-D array
$A \in \R^{m \times n}$	A.shape == (m, n)	$m$ rows, $n$ columns
$x_i$, $i = 1, \dots, n$	x[i - 1]	math counts from 1, Python from 0
$A_{ij}$	A[i - 1, j - 1]	row first, then column
$\sum_{i=1}^n x_i$	np.sum(x) or x.sum()
$\sum_i w_i x_i = w^\top x$	np.sum(w * x) or w @ x
$Ax$, $AB$	A @ x, A @ B	matrix product
$A^\top$	A.T	transpose

Indexing selects cells. In the figure, A = np.arange(30).reshape(5, 6) holds the numbers 0 to 29, row by row. An index for the rows and an index for the columns pick cells. An integer picks one row (or column) and removes that axis; a slice start:stop:step picks several and keeps the axis.

Interactive figure: loading…

A reduction collapses one axis. X.sum(axis=0) adds down the rows: each column becomes one number. X.sum(axis=1) adds along each row: each row becomes one number. Read axis=k as "the axis with index $k$ disappears". The shape (3, 4) becomes (4,) for axis=0 and (3,) for axis=1. With keepdims=True the axis stays, with length 1.

Interactive figure: loading…

Broadcasting stretches size-1 dimensions. NumPy can add arrays of different shapes. It aligns the shapes from the right. In each position the two sizes must be equal, or one of them must be 1 (a missing dimension counts as 1). A size-1 dimension is then repeated ("stretched") to match the other.

Interactive figure: loading…

3. Definitions and main results #

3.1 Python in ten minutes #

This is the Python that this site uses. Type each line in the notebook and compare with the comment.

x = 3              # an int (integer)
y = 2.5            # a float (a number with a decimal point)
print(x / 2, x // 2, x ** 2)    # 1.5 1 9   (/ divides, // divides and rounds down, ** is a power)

xs = [1, 2, 3]     # a list: an ordered collection
xs.append(4)       # add an element at the end
print(xs[0], xs[-1], len(xs))   # 1 4 4    (index 0 is the first element, -1 the last)

total = 0
for v in xs:       # a for loop runs the indented block once per element
    total += v     # the same as total = total + v
print(total)       # 10

def square(t):     # a function: def, name, arguments, colon, indented body
    return t * t

squares = [square(v) for v in xs]       # a list comprehension: [1, 4, 9, 16]
print(f"sum = {total}, mean = {total / len(xs):.2f}")   # an f-string: sum = 10, mean = 2.50

import numpy as np  # load NumPy under the short name np (everyone uses this name)

Indentation (4 spaces) is part of the syntax: it shows which lines belong to the loop or the function. Lists are not vectors: [1, 2] * 2 is [1, 2, 1, 2] (repetition) and [1, 2] + [3, 4] is [1, 2, 3, 4] (concatenation). For math, use NumPy arrays: np.array([1, 2]) * 2 is [2 4]. The official Python tutorial has more.

3.2 The ndarray: shape, dtype, ndim #

A NumPy array (ndarray, "n-dimensional array") holds numbers of one type in a grid.

• The shape is a tuple with the length of each axis: (3,) for 3 numbers in a line, (2, 3) for 2 rows and 3 columns. Note the comma in (3,): it is a tuple with one element.
• ndim is the number of axes: len(a.shape). An axis is one direction of the grid; axis 0 is the first index (rows of a 2-D array), axis 1 the second (columns).
• The dtype is the type of the entries, for example float64 (a 64-bit float, the default for decimals) or int64 (a 64-bit integer).

a = np.array([1.0, 2.0, 3.0])
print(a.shape, a.ndim, a.dtype)            # (3,) 1 float64
A = np.array([[1, 2, 3], [4, 5, 6]])
print(A.shape, A.ndim, A.dtype)            # (2, 3) 2 int64

The integer type depends on the platform: int64 on Colab and most computers, int32 in the in-browser notebooks and older Windows installations. Use floats for math unless you need integers.

3.3 Creating arrays #

Code	Result	Note
np.zeros((2, 3))	a $2 \times 3$ array of 0.	the shape is one tuple argument
np.ones(3)	[1. 1. 1.]
np.eye(3)	the $3 \times 3$ identity matrix $I$
np.full((2, 2), 7.0)	a $2 \times 2$ array of 7.
np.arange(0, 1, 0.25)	[0. 0.25 0.5 0.75]	start, stop, step; stop is excluded
np.linspace(0, 1, 5)	[0. 0.25 0.5 0.75 1.]	start, stop, number of points; stop is included

np.linspace(a, b, n) gives $n$ points with step $(b - a)/(n - 1)$. Use it for plotting grids. Use np.arange for integer counters (np.arange(5) is [0 1 2 3 4]). With a float step, rounding can change the number of points: np.arange(1, 1.3, 0.1) returns 4 numbers, [1. 1.1 1.2 1.3]. NumPy computes the number of points as $\lceil (\text{stop} - \text{start})/\text{step} \rceil$, and in floating point $(1.3 - 1)/0.1$ is $3.0000000000000004$, which rounds up to 4. The last point is about $1.3000000000000003$, slightly larger than the stop.

3.4 Indexing, slicing, masks; views and copies #

Python counts from 0. For an array x of length $n$:

• x[0] is the first element ($x_1$ in math), x[n - 1] or x[-1] the last.
• x[i:j] is the slice of elements i, i+1, …, j-1: j is excluded, and the length is j - i (for $0 \le i \le j \le n$). x[:j] starts at 0, x[i:] runs to the end, x[::2] takes every second element, x[::-1] reverses.
• For a 2-D array A of shape (m, n), A[i, j] is one entry, A[i] (or A[i, :]) is row i, and A[:, j] is column j. Both A[i] and A[:, j] are 1-D, with shape (n,) and (m,).
• A boolean mask selects by a condition: x > 0 is an array of True/False, and x[x > 0] keeps the entries where it is True. For m = np.array([3, -1, 4, -1, 5]), m[m > 0] is [3 4 5].

Rule (integer versus slice). An integer index removes its axis; a slice keeps it, even when the slice has length 1. So for A of shape (5, 6), A[1] has shape (6,) and A[1:2] has shape (1, 6).

Views and copies. A slice does not copy the numbers. It is a view: a new array object that looks at the same memory. Writing into a view changes the original. A boolean mask (and an index list such as x[[0, 2]]) returns a copy. When you need an independent array, call .copy().

a = np.arange(5)        # [0 1 2 3 4]
b = a[1:3]              # a view of a[1] and a[2]
b[0] = 99
print(a)                # [ 0 99  2  3  4]   a changed too
c = a[1:3].copy()       # an independent copy
c[0] = -5
print(a)                # [ 0 99  2  3  4]   a did not change

NumPy uses views for speed: slicing a large array costs no memory. The price is that you must know when two names share the same numbers.

3.5 Elementwise operations and reductions #

Elementwise operations act on each entry: x + y, x - y, x * y, x / y, x ** 2, np.exp(x), np.log(x), np.sqrt(x), np.abs(x). For two arrays of the same shape, entry i of x * y is x[i] * y[i]. A number combines with every entry: 2 * x, x + 1.

A reduction turns many numbers into fewer: sum, mean, std, min, max, argmin, argmax. Without axis it uses all entries and returns one number. With axis=k it works along axis k only.

Rule (shape of a reduction). For an array of shape $(d_0, d_1, \dots)$, a reduction with axis=k returns shape $(d_0, \dots, d_{k-1}, d_{k+1}, \dots)$: axis $k$ is removed. With keepdims=True, $d_k$ is replaced by 1 instead.

In math: for $X \in \R^{m \times n}$, X.sum(axis=0) is the vector with entries $\sum_{i=1}^m X_{ij}$ (one per column $j$), and X.sum(axis=1) has entries $\sum_{j=1}^n X_{ij}$ (one per row $i$). For a data matrix with samples as rows, X.mean(axis=0) is the mean of each feature. np.argmax(x) is the position of the largest entry (the first one if there is a tie), not the largest value; x.max() is the value.

3.6 Shapes: reshape, transpose, new axes #

• a.reshape(m, n) gives the same numbers in a new shape, filled row by row; $m \cdot n$ must equal the number of entries. One size may be -1: NumPy computes it. a.reshape(-1) flattens to 1-D.
• A.T is the transpose: shape (m, n) becomes (n, m).
• x[:, None] (or x[:, np.newaxis]) adds an axis of length 1: shape (n,) becomes (n, 1), a column. x[None, :] gives (1, n), a row.

A 1-D array has no orientation. Shape (n,) is neither a row nor a column, and x.T does nothing to it. NumPy prints it on one line. In math we treat vectors as columns; in NumPy a 1-D array behaves like a column on the right of @ (A @ x) and like a row on the left (x @ A). When you need a real column, make it: x[:, None] or x.reshape(-1, 1).

x = np.arange(3)                                  # [0 1 2]
print(x.shape, x.T.shape)                         # (3,) (3,)
print(x[:, None].shape, x[None, :].shape)         # (3, 1) (1, 3)
print(np.arange(6).reshape(2, 3))                 # [[0 1 2]
                                                  #  [3 4 5]]

3.7 Elementwise product * versus matrix product @ #

* is the elementwise product (also called the Hadamard product, written $A \odot B$ in some papers): $(A \odot B)_{ij} = A_{ij} B_{ij}$. It needs equal shapes, or shapes that broadcast.

@ is the matrix product $AB$: $(AB)_{ij} = \sum_k A_{ik} B_{kj}$, row $i$ of $A$ dotted with column $j$ of $B$. Shapes (m, p) @ (p, n) give (m, n): the inner sizes must be equal. Linear Algebra lesson 1 explains why the product is defined this way. For two 1-D arrays of the same length, x @ y is the dot product $x^\top y = \sum_i x_i y_i$, one number.

Interactive figure: loading…

3.8 Broadcasting: the rule #

Rule (broadcasting). To combine arrays with shapes $s$ and $t$ in an elementwise operation:

1. Write the two shapes one below the other, aligned on the right. If one shape is shorter, add 1s on its left until the lengths are equal.
2. Compare the sizes in each position. They are compatible if they are equal or if one of them is 1. If some position is not compatible, NumPy raises ValueError: operands could not be broadcast together with shapes ….
3. The result size in each position is the larger of the two sizes. An array with size 1 in a position behaves as if its values were repeated along that axis. NumPy does not copy the data to do this.

a	b	aligned	result
(3, 4)	(4,)	(3, 4) and (1, 4)	(3, 4): b is added to every row
(3, 1)	(4,)	(3, 1) and (1, 4)	(3, 4): the outer sum a[i, 0] + b[j]
(3,)	(3, 1)	(1, 3) and (3, 1)	(3, 3): often a bug, see below
(3, 4)	(3, 1)	(3, 4) and (3, 1)	(3, 4): b[i, 0] is added to row i
(3, 4)	(3,)	(3, 4) and (1, 3)	error: 4 and 3 in the last position
(2, 3)	(2,)	(2, 3) and (1, 2)	error: 3 and 2 in the last position

A number (a scalar) has shape () and broadcasts against everything: x + 1 adds 1 to every entry.

Centering data, step by step. Let $X \in \R^{m \times n}$ hold $m$ samples as rows. The column means are $\mu_j = \frac{1}{m} \sum_{i=1}^m X_{ij}$, so mu = X.mean(axis=0) has shape (n,). Then X - mu has shape (m, n) and entry $X_{ij} - \mu_j$: mu is aligned with the last axis (the columns) and repeated for every row. Each centered column has mean 0:

To center the rows instead, the means must keep their axis: X - X.mean(axis=1, keepdims=True) subtracts shape (m, 1) from (m, n). Without keepdims, the means have shape (m,), which is aligned with the columns: you get an error when $m \ne n$, and a silent wrong answer when $m = n$.

The silent (n,) versus (n, 1) bug. Let y have shape (n,) and y_hat shape (n, 1) (a column, for example from X @ w with w of shape (d, 1)). Then y - y_hat has shape (n, n), not (n,), and no error is raised. Its entry [i, j] is y[j] - y_hat[i, 0]: every target minus every prediction. The mean of its squares is a number, but it is not the mean squared error. Section 4 computes both. The fix: make the shapes equal (y_hat.ravel() or y_hat[:, 0]), and print shapes.

Why NumPy chose this rule

Aligning on the right means the last axis is matched first. In a data matrix the last axis is the features, so a vector of per-feature numbers (means, scales, weights) broadcasts across all samples without any reshaping. In an image batch (B, C, H, W), a per-channel mean has shape (C, 1, 1): the two 1s line it up with H and W, so it is stretched over every pixel and every image. The same rule is used by PyTorch, JAX and TensorFlow, so it is worth learning once.

3.9 Vectorization: from Σ to one line #

Vectorization means writing an operation on whole arrays instead of a Python loop over entries. The loop then runs inside NumPy's compiled code, which is much faster. The notebook times both versions; the speed-up depends on the computer and the size, so measure it.

Math	Loop	Vectorized
$\sum_{i=1}^n w_i x_i$	sum(w[i] * x[i] for i in range(n))	w @ x or np.sum(w * x)
$\frac{1}{n}\sum_{i=1}^n (y_i - \hat y_i)^2$ (MSE)	a loop with a running total	np.mean((y - y_hat) ** 2)
$\hat y^{(i)} = w^\top x^{(i)} + b$ for all samples	a loop over rows of X	X @ w + b
$\|x\| = \sqrt{\sum_i x_i^2}$	a loop and math.sqrt	np.sqrt(np.sum(x ** 2)) or np.linalg.norm(x)

In X @ w + b, X has shape (m, n), w has shape (n,), so X @ w has shape (m,): one prediction per sample ($x^{(i)}$ is row $i$ of $X$, as in lesson 02). The number b broadcasts to every entry.

Result (pairwise squared distances). For vectors $a, b \in \R^n$,

Proof.

1. By definition $\|v\|^2 = v^\top v$, so $\|a - b\|^2 = (a - b)^\top (a - b)$.
2. Multiply out: $(a - b)^\top (a - b) = a^\top a - a^\top b - b^\top a + b^\top b$.
3. $b^\top a = \sum_i b_i a_i = \sum_i a_i b_i = a^\top b$, so the two middle terms add to $-2a^\top b$.
4. With $a^\top a = \|a\|^2$ and $b^\top b = \|b\|^2$ we get the formula. $\blacksquare$

For all pairs of rows of $P \in \R^{m \times n}$ at once: sq = (P ** 2).sum(axis=1) has shape (m,); G = P @ P.T has shape (m, m) and $G_{ij} = p_i^\top p_j$ (rows $p_i$, $p_j$). Then D = sq[:, None] + sq[None, :] - 2 * G has shape (m, m) and $D_{ij} = \|p_i - p_j\|^2$. The broadcast (m, 1) + (1, m) builds the table of all sums $\|p_i\|^2 + \|p_j\|^2$. With floats, tiny negative values such as -1e-15 can appear on the diagonal; np.maximum(D, 0) removes them. This is how k-nearest neighbours, k-means and RBF kernels compute distances.

3.10 Plotting with matplotlib #

matplotlib draws plots. The standard pattern: make a figure and axes, draw on the axes, add labels, show.

import matplotlib.pyplot as plt

x = np.linspace(-2, 2, 200)                 # 200 points: a smooth curve
fig, ax = plt.subplots(figsize=(5, 3))      # one figure with one set of axes
ax.plot(x, x ** 2, label="$x^2$")           # a line through the points (x[i], x[i]**2)
ax.plot(x, np.exp(x), label="$e^x$")
ax.set_xlabel("x"); ax.set_ylabel("y")
ax.set_title("Two functions"); ax.legend()
plt.show()

ax.plot joins the points with straight lines, so a grid of 200 points looks smooth; with 5 points you would see corners. ax.scatter(x, y) draws points without lines, and ax.hist(data, bins=30) draws a histogram. fig, axes = plt.subplots(1, 3, figsize=(12, 3)) makes three plots side by side; axes[0] is the first one. Lesson 04 uses these plots for exp and log.

3.11 Random numbers #

ML uses random numbers everywhere: weight initialization, shuffling, dropout, train/test splits, Monte Carlo estimates. NumPy's modern interface is a generator object:

rng = np.random.default_rng(0)          # a generator with seed 0
print(rng.normal(size=3))               # [ 0.1257 -0.1321  0.6404]  (rounded)
print(rng.uniform(0, 1, size=3))        # 3 numbers, uniform on [0, 1)
print(rng.integers(0, 6, size=10))      # 10 integers from 0 to 5: high is excluded
print(rng.choice([10, 20, 30, 40], size=2, replace=False))   # 2 different elements

The numbers are pseudo-random: a deterministic algorithm produces them from a starting value, the seed. The same seed gives the same sequence, on every run. This is reproducibility: you, a reviewer or a test get exactly the same result. Without a seed (np.random.default_rng()), the seed comes from the operating system and every run differs.

Interactive figure: loading…

You meet the normal distribution, histograms and the sample mean properly in lesson 08. Older code uses np.random.seed(0) and np.random.rand(...); this works, but new code should use a generator.

3.12 Floats and integers: four surprises #

1. Decimals are not exact. A float stores a binary fraction, and $0.1$ has no exact binary form. So 0.1 + 0.2 is 0.30000000000000004, and 0.1 + 0.2 == 0.3 is False. Compare floats with a tolerance: np.isclose(a, b) checks $|a - b| \le 10^{-8} + 10^{-5}|b|$; np.allclose(A, B) does this for every entry and returns one True/False.
2. Overflow gives inf. The largest float64 is about $1.8 \times 10^{308}$. np.exp(709) is about $8.2 \times 10^{307}$, but np.exp(1000) is inf, with a RuntimeWarning: overflow. Lesson 05 shows how softmax avoids this.
3. Integer arrays stay integer. a = np.array([1, 2, 3]) has an integer dtype. a[0] = 2.7 stores 2: the decimal part is cut off without a warning. Create float arrays (np.array([1., 2., 3.])) when you will store decimals.
4. Two divisions. / always gives a float (7 / 2 is 3.5). // divides and rounds down (7 // 2 is 3, and -7 // 2 is -4, not -3).

Linear Algebra lesson 13 explains floating point in full.

Answers to section 1. (1) (3,): axis 0 (the 100 samples) disappears, and you get one mean per feature. (2) Yes, and the result has shape (3, 3): aligned from the right, (3,) acts as (1, 3), and both size-1 dimensions are stretched. This is the silent bug of section 3.8. (3) [ 0 99 2 3 4]: b is a view, so writing into b[0] writes into a[1].

4. Worked example #

(a) Standardize a small data matrix. Four samples, two features:

Step 1: the column means, mu = X.mean(axis=0), shape (2,). $\mu_1 = (1 + 3 + 3 + 1)/4 = 8/4 = 2$ and $\mu_2 = (13 + 21 + 27 + 19)/4 = 80/4 = 20$. So mu is [ 2. 20.].

Step 2: center, Xc = X - mu. Shapes (4, 2) and (2,) align as (4, 2) and (1, 2), so mu is subtracted from every row:

Each column of $X_c$ adds up to 0, as section 3.8 proved.

Step 3: the standard deviations, sd = X.std(axis=0) (NumPy divides by $m$; lesson 08 discusses dividing by $m - 1$). Column 1: $\sqrt{(1 + 1 + 1 + 1)/4} = \sqrt{1} = 1$. Column 2: $\sqrt{(49 + 1 + 49 + 1)/4} = \sqrt{25} = 5$. So sd is [1. 5.].

Step 4: scale, Z = Xc / sd, again (4, 2) with (2,): column 1 is divided by 1, column 2 by 5.

Each column of $Z$ has mean 0 and standard deviation 1: column 2 has mean $(-1.4 + 0.2 + 1.4 - 0.2)/4 = 0$ and $\sqrt{(1.96 + 0.04 + 1.96 + 0.04)/4} = \sqrt{1} = 1$. In code the mean of column 2 prints as about $-1.4 \times 10^{-17}$, not exactly 0: a float rounding error, so check with np.allclose.

(b) Pairwise squared distances. Three points in the plane, as the rows of $P$:

Step 1: squared norms of the rows, sq = (P ** 2).sum(axis=1) = $(0,\ 9 + 16,\ 1 + 1) = (0, 25, 2)$, shape (3,).

Step 2: the dot products of all pairs of rows, G = P @ P.T, shape (3, 3). For example $G_{23} = (3, 4)^\top (1, 1) = 3 + 4 = 7$ (math indices from 1):

Step 3: sq[:, None] has shape (3, 1) and sq[None, :] shape (1, 3), so their sum has shape (3, 3) with entry $\|p_i\|^2 + \|p_j\|^2$. Subtract $2G$:

Check one entry by hand: $\|p_2 - p_3\|^2 = (3 - 1)^2 + (4 - 1)^2 = 4 + 9 = 13$ ✓.

(c) The silent shape bug. Targets y = [1, 2, 3] with shape (3,), predictions y_hat = [[1], [2], [4]] with shape (3, 1).

• Correct MSE: $\frac{1}{3}\big((1-1)^2 + (2-2)^2 + (3-4)^2\big) = \frac{1}{3}$.
• np.mean((y - y_hat) ** 2): y - y_hat has shape (3, 3), entry [i, j] = y[j] - y_hat[i, 0]:

whose mean is $21/9 \approx 2.333$. No error, a wrong number. The fix is y_hat[:, 0] or y_hat.ravel().

5. In code #

NumPy checks the hand computations above. The notebook has these cells, more demos (a timing of loop versus vectorized code, plots) and the exercises.

import numpy as np
from scipy.spatial.distance import cdist

# (a) standardize with broadcasting
X = np.array([[1., 13.], [3., 21.], [3., 27.], [1., 19.]])
mu, sd = X.mean(axis=0), X.std(axis=0)
print(X.shape, mu.shape, mu, sd)        # (4, 2) (2,) [ 2. 20.] [1. 5.]
Z = (X - mu) / sd
print(Z[:, 1])                          # [-1.4  0.2  1.4 -0.2]
print(np.allclose(Z.mean(axis=0), 0), np.allclose(Z.std(axis=0), 1))   # True True

# (b) pairwise squared distances without loops
P = np.array([[0., 0.], [3., 4.], [1., 1.]])
sq = (P ** 2).sum(axis=1)               # [ 0. 25.  2.]
D = sq[:, None] + sq[None, :] - 2 * P @ P.T
print(D.tolist())                       # [[0.0, 25.0, 2.0], [25.0, 0.0, 13.0], [2.0, 13.0, 0.0]]
print(np.allclose(D, cdist(P, P, "sqeuclidean")))   # True

# (c) the silent shape bug
y = np.array([1., 2., 3.])
y_hat = np.array([[1.], [2.], [4.]])
print((y - y_hat).shape, np.mean((y - y_hat) ** 2))     # (3, 3) 2.333...
print(np.mean((y - y_hat[:, 0]) ** 2))                  # 0.333...

# * versus @, and a seeded generator
x = np.array([1, 2, 3])
print(x * x, x @ x)                     # [1 4 9] 14
rng = np.random.default_rng(0)
a = rng.normal(size=3)
b = np.random.default_rng(0).normal(size=3)
print(np.array_equal(a, b))             # True: same seed, same numbers
print(0.1 + 0.2 == 0.3, np.isclose(0.1 + 0.2, 0.3))     # False True

In Python, * and @ have the same precedence and run from left to right, so 2 * P @ P.T means (2 * P) @ P.T. This equals 2 * (P @ P.T), so the result is the same.

6. Where it shows up #

• Data preprocessing. scikit-learn's StandardScaler computes (X - X.mean(axis=0)) / X.std(axis=0), exactly section 4(a). Image pipelines subtract a per-channel mean of shape (C, 1, 1) from images of shape (C, H, W).
• Neural networks. A linear layer on a batch is X @ W + b: X has shape (batch, d_in), W shape (d_in, d_out), and the bias b of shape (d_out,) broadcasts over the batch. The predicted classes are logits.argmax(axis=1); the accuracy is np.mean(pred == labels).
• Attention. In a transformer, the scores Q @ K.T / np.sqrt(d_k) form a matrix of shape (T, T), a softmax runs along axis=-1 (each row), and a causal mask is added by broadcasting.
• Nearest neighbours and kernels. k-nearest neighbours, k-means and the RBF kernel $\exp(-\gamma \|x - x'\|^2)$ all start from the pairwise distance trick of section 3.9.
• Reproducible experiments. Papers and libraries fix seeds (np.random.default_rng(seed), torch.manual_seed(seed)) so that results can be repeated and compared.
• ML for science. Simulation outputs are arrays such as (time, lat, lon) in climate data or (frames, atoms, 3) in molecular dynamics. "Mean over time" is axis=0, and anomalies are data - data.mean(axis=0): broadcasting again.

7. Common mistakes #

• Off-by-one indices. Math $x_1$ is x[0], and x[1:3] has 2 elements, not 3. The stop index is excluded in slices, in range and in np.arange.
• The wrong axis. X.mean(axis=1) on a data matrix averages over the features of each sample, not over samples. Say "axis 0 disappears" aloud and check the shape.
• (n,) mixed with (n, 1). y - y_hat becomes (n, n) without an error. Print shapes, and add assert y_hat.shape == y.shape in your code.
• * where you meant @. For square matrices A * B runs without error and gives a different matrix than A @ B. Read * as "entry by entry".
• Writing into a slice. b = a[1:3] and then changing b changes a. Use .copy() when you need an independent array.
• Comparing floats with ==. Use np.isclose or np.allclose.
• Storing decimals in an integer array. np.zeros(3, dtype=int) or np.array([1, 2, 3]) cuts off decimals on assignment. Use float arrays for math.
• Forgetting the seed, then a result cannot be repeated. Create one rng at the top of the notebook and pass it around.

8. Check yourself #

1. A = np.arange(12).reshape(3, 4). Give A[1], A[:, -1] and A[1:, :2], with their shapes.
   Answer

   `A` has rows `[0 1 2 3]`, `[4 5 6 7]`, `[8 9 10 11]`. `A[1]` is `[4 5 6 7]`, shape `(4,)`. `A[:, -1]` is the last column `[ 3 7 11]`, shape `(3,)`. `A[1:, :2]` is rows 1–2 and columns 0–1: `[[4 5] [8 9]]`, shape `(2, 2)`.
2. What are the shapes of A.sum(), A.sum(axis=0), A.sum(axis=1, keepdims=True) for the same A?
   Answer

   `()` (one number, 66), `(4,)` (one sum per column: `[12 15 18 21]`), and `(3, 1)` (one sum per row, as a column: `[[6] [22] [38]]`).
3. Which of these work, and what is the result shape: (5, 3) + (3,), (5, 3) + (5,), (5, 3) + (5, 1), (4, 1, 3) + (2, 1)?
   Answer

   `(5, 3)`; error (3 and 5 meet in the last position); `(5, 3)`; `(4, 2, 3)` (align `(4, 1, 3)` with `(1, 2, 1)`: every position has equal sizes or a 1).
4. Write one line that computes $\frac{1}{n}\sum_{i=1}^n |y_i - \hat y_i|$ (the mean absolute error) for 1-D arrays y and y_hat.
   Answer

   `np.mean(np.abs(y - y_hat))`. First check that both have shape `(n,)`.
5. rng = np.random.default_rng(1); a = rng.normal(size=2); b = rng.normal(size=2). Is a equal to b? What would make them equal?
   Answer

   No. `b` continues the same sequence, so it holds the next two numbers. To get equal arrays, create a new generator with the same seed before each draw: `np.random.default_rng(1).normal(size=2)` twice.
6. Why does np.arange(0, 1, 0.1) have 10 elements but np.linspace(0, 1, 10) does not contain 0.1?
   Answer

   `arange` stops before 1 with step 0.1: $0, 0.1, \dots, 0.9$ (10 numbers). `linspace(0, 1, 10)` includes both ends, so its step is $1/9 \approx 0.111$: the points are $0, 0.111, \dots, 1$. For steps of 0.1 with the end included, use `np.linspace(0, 1, 11)`.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. State the broadcasting rule of NumPy.

   Show answer

2. X has shape (m, n). What is the shape of X.mean(axis=0), and what does it contain?

   Show answer

3. What is the shape of a + b for a.shape == (3, 1) and b.shape == (4,)?

   Show answer

4. What is the difference between A * B and A @ B?

   Show answer

5. b = a[1:3], then b[0] = 99. Why does a change?

   Show answer

6. How do np.arange(0, 1, 0.25) and np.linspace(0, 1, 5) differ?

   Show answer

7. How do you make NumPy random numbers reproducible?

   Show answer

8. Why should you not test floats with ==?

   Show answer

9. What is the difference between shapes (n,), (n, 1) and (1, n)?

   Show answer

10. Write the pairwise squared distances of the rows of P without loops.

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "What is broadcasting? Give a case where it hides a bug." — Shapes align on the right; equal sizes or 1. The classic bug: y of shape (n,) minus predictions of shape (n, 1) gives (n, n), and the loss is wrong without an error. Fix by matching shapes and asserting them.
• "Standardize the columns of X without a loop." — (X - X.mean(axis=0)) / X.std(axis=0); explain that the (n,) statistics broadcast over the rows. Mention a small epsilon or a check for constant columns (std 0).
• "Compute all pairwise Euclidean distances between two sets of points." — For A of shape (m, d) and B of shape (k, d): (A**2).sum(1)[:, None] + (B**2).sum(1)[None, :] - 2 * A @ B.T, clip at 0, then np.sqrt. Memory is $O(mk)$ instead of $O(mkd)$ for the naive A[:, None, :] - B[None, :, :].
• "Why is vectorized code faster than a Python loop?" — The loop runs in compiled C code over contiguous memory, without the per-element overhead of the Python interpreter, and it can use SIMD instructions and optimized BLAS libraries for @.
• "View or copy?" — Basic slicing returns a view; fancy indexing (index arrays, boolean masks) returns a copy. np.shares_memory(a, b) tells you.

10. Practice #

The notebook for this lesson runs in Colab or in your browser. It has five exercises, each with a check cell and a folded solution:

1. Write the dot product $\sum_i w_i x_i$ three ways (a loop, np.sum(w * x), w @ x) and check that they agree.
2. Standardize a data matrix with broadcasting so that every column has mean 0 and standard deviation 1.
3. Compute the pairwise squared distances between the rows of X without loops, and compare with scipy.spatial.distance.cdist.
4. Make the arrays for a plot of $\sin$ and $\cos$ on $[0, 2\pi]$, then plot them with labels and a legend.
5. Estimate $\pi$ by Monte Carlo with a seeded generator, and check that the same seed gives the same estimate.

11. Watch and read #

Complete Python NumPy Tutorial (Creating Arrays, Indexing, Math, Statistics, Reshaping)Keith Galli · 59 min

A hands-on tour of array creation, indexing, math, statistics and reshaping; type along in the notebook.

Numpy Array Broadcasting In Python ExplainedmCoding · 13 min

The broadcasting rule with pictures, including the right-alignment of shapes from section 3.8.

Matplotlib Tutorial (Part 1): Creating and Customizing Our First PlotsCorey Schafer · 35 min

Line plots, labels, legends and styles; the same plt.subplots pattern as section 3.10.

• Stanford CS231n, Python Numpy Tutorial: Python, NumPy arrays, indexing, broadcasting and matplotlib on one page, written for ML students.
• NumPy documentation, NumPy: the absolute basics for beginners and Broadcasting (the rule of section 3.8, with figures).
• NumPy documentation, Random Generator: default_rng and its methods.
• Jake VanderPlas, Python Data Science Handbook (free online), chapter 2 (NumPy) and chapter 4 (matplotlib).
• The official Python tutorial, sections 3–5, for more Python.