Python Basics for ML

Python Basics for ML

Why it matters

ML code runs on arrays, and the single biggest habit to build is vectorisation: expressing a computation as one operation on whole arrays instead of a Python loop over elements. Summing a million floats with a Python loop takes about a hundred times longer than arr.sum(). Every forward pass, loss and gradient you write later is built from the three operations in this section: reductions, matrix multiplication and broadcasting.

Core ideas

Vectorise, don't loop

A Python loop pays interpreter overhead for every element. A NumPy call runs one compiled loop over contiguous memory, often with SIMD instructions. Rule of thumb: if you are writing for i in range(len(arr)), there is almost always an array expression that does the same thing.

Reductions

A reduction collapses values into fewer values: arr.sum(), arr.mean(), arr.max(). With no axis it reduces everything to a scalar; axis=0 reduces down the rows (one result per column) and axis=1 across the columns (one result per row).

Matrix multiplication

For A of shape (m, n) and B of shape (n, p):

C[i, j] = sum over k of A[i, k] * B[k, j]        C has shape (m, p)

The inner dimensions (n) must match and disappear. It costs m * n * p multiply-adds. Three spellings:

Call	2D inputs	Batched 3D inputs
A @ B / np.matmul(A, B)	matrix product	batched product, (b, m, n) @ (b, n, p) -> (b, m, p)
np.dot(A, B)	matrix product	a different contraction: (2, 3, 4) dot (2, 4, 5) gives (2, 3, 2, 5)
A * B	elementwise product, not matmul	elementwise

Prefer @: it reads like the maths and does the right thing for batches.

Broadcasting

When shapes differ, NumPy compares them from the right; dimensions must be equal or 1, and size-1 (or missing) dimensions are stretched. That is how a bias vector of shape (n,) is added to every row of an (m, n) matrix without copying it m times.

Worked example: a linear layer by hand

A linear layer computes y = x @ W + b. With a batch of 2 samples, 2 input features and 2 outputs:

x = np.array([[1, 2],
              [3, 4]])           # (batch=2, d_in=2)
W = np.array([[5, 6],
              [7, 8]])           # (d_in=2, d_out=2)
b = np.array([10, 20])           # (d_out,)

x @ W        # [[19, 22], [43, 50]]   e.g. 19 = 1*5 + 2*7
x @ W + b    # [[29, 42], [53, 70]]   b broadcasts over the batch
(x @ W).sum()  # 134                  a reduction to one number

All three operations of this section show up in one line of a neural network.

Common pitfalls

• * is not matrix multiplication: A * B multiplies elementwise (and broadcasts), so it can silently "work" with the wrong meaning.
• Shape mismatches: (2, 3) @ (2, 3) raises a ValueError (3 vs 2 in the inner dimension). Transpose the operand you meant: A @ B.T.
• 1D vectors are neither rows nor columns: v @ A treats v as a row and A @ v as a column, and both return 1D. Reshape to (n, 1) when you need an explicit column.
• Integer overflow: small integer dtypes wrap around (two int8 values of 100 add up to -56). Use float arrays for numeric work.
• Floats aren't exact: 0.1 + 0.2 == 0.3 is False. Compare floats with a tolerance (np.isclose, np.allclose).

Interview angle

• "Implement matrix multiply" usually means: state the shape rule, write A @ B, then possibly the triple loop, and give its O(m·n·p) cost.
• Follow-ups: why is the vectorised version faster; what's the difference between np.dot, np.matmul and *; how would you multiply a batch of matrices; what happens with mismatched shapes.
• Be ready to say how broadcasting adds a bias to a batch, and what shape the result has.

Problems in this section

1. Array Sum: reduce an array to a single number with np.sum.
2. Matrix Multiplication: multiply (m, n) by (n, p) with @, and check the shape rule.
3. Broadcasting Addition: add a vector to every row of a matrix with one +.

For deeper practice, see NumPy Fundamentals (indexing, axes, reshape) and Einstein Summation, which writes matrix products as index expressions like 'ik,kj->ij'.