NumPy · Module 10 of 13 · 6 concepts · about 8 min

Linear Algebra with NumPy

Matrix products, norms and pairwise scores — the four operations behind every similarity search.

At work: The maths under recommendations, search and machine learning: matching users to products, documents to queries.

matmul versus elementwise multiply

What it is

A @ B (equivalently np.matmul) is the matrix product: the last axis of A contracts with the second-to-last of B. A * B is element-wise and broadcasts.

Why it matters

They are different operations that accept overlapping shapes, so confusing them is a wrong answer rather than an error — A * B on two (n, n) matrices runs happily and means nothing you wanted.

In the real world: Multiplying a customers-by-products purchase matrix by a products-by-price vector gives each customer's total spend in one step.

Visual example

matmul

A matrix product contracts the inner axes: (2, 3) @ (3, 2) gives (2, 2), one dot product per output cell.

A — shape (2, 3)
012
0102
1011
A @ B
B — shape (3, 2)
01
012
130
211
A @ B — shape (2, 2)
01
034
141

Cell [0, 0], and its neighbours

[0,0] = 1*1 + 0*3 + 2*1 = 3

[0,1] = 1*2 + 0*0 + 2*1 = 4

[1,0] = 0*1 + 1*3 + 1*1 = 4

[1,1] = 0*2 + 1*0 + 1*1 = 1

Why this matters

Every scoring step — queries against documents, activations against weights — is one matmul, handed to a tuned BLAS kernel that no hand-written loop will approach.

Common misconception

Wrong idea: A * B multiplies matrices.

Correct idea: * is element-wise; @ is the matrix product. The shape rule differs too: @ needs the inner axes to agree, not the whole shape.

  1. 1Check the inner axes(2, 3) @ (3, 2): the 3s must agree and then vanish.
  2. 2One dot product per cellOutput [i, j] is row i of A against column j of B.
  3. 3Read the shapeWhat is left is the outer axes: (2, 2).
A = np.arange(6).reshape(2, 3)
B = np.arange(6).reshape(3, 2)
(A @ B).shape        # (2, 2) — (2, 3) x (3, 2)
(A * A).shape        # (2, 3) — element-wise
A @ A                # raises: (2, 3) x (2, 3) does not contract

Using np.dot on arrays with more than two axes. For 2-D it matches @, but for higher dimensions it contracts differently from matmul, and the difference appears as a batch axis in the wrong place. Use @ and mean it.

Batched matmul and broadcasting

What it is

@ treats all but the last two axes as batch axes and broadcasts them: (b, n, d) @ (d, m) gives (b, n, m), and (b, n, d) @ (b, d, m) multiplies each batch element separately.

Why it matters

It is how attention and any per-example transform are written without a loop over the batch, and it is the shape convention every deep learning framework inherited.

In the real world: Scoring many users against many products at once, the core step of a recommendation engine.

Q = np.zeros((32, 10, 64))      # batch of 32, 10 tokens, 64 dims
W = np.zeros((64, 64))
(Q @ W).shape                   # (32, 10, 64) — W is shared across the batch
(Q @ Q.transpose(0, 2, 1)).shape  # (32, 10, 10) — per-example token scores

Transposing a batched array with .T, which reverses *every* axis and puts the batch axis last. Use transpose(0, 2, 1), or swapaxes(-1, -2), which says what it means at any ndim.

Practise it: Clip Outliers to a Percentile Range, Cosine Similarity Without Loops, Normalise Rows to Unit Length, One-Hot Encode a Label Vector, Pairwise Euclidean Distances, Pick the Rows That Pass a Threshold, Scale Every Column to 0 and 1, Intersection Over Union for Every Box Pair, Softmax That Does Not Overflow

Norms with np.linalg.norm

What it is

np.linalg.norm(x, axis=, keepdims=) computes vector or matrix norms. With axis=1 on an (n, d) matrix it gives the length of each row; ord=1 gives the sum of absolute values, ord=np.inf the maximum.

Why it matters

Row lengths are the denominator of cosine similarity and the scale in gradient clipping, and forgetting the axis argument collapses the whole matrix into one number instead.

In the real world: Measuring how far each product's embedding is from the origin, a step before comparing products by similarity.

X = np.array([[3.0, 4.0], [0.0, 0.0]])
np.linalg.norm(X, axis=1)                 # array([5., 0.])
np.linalg.norm(X, axis=1, keepdims=True)  # shape (2, 1), ready to divide by
np.linalg.norm(X)                          # 5.0 — the Frobenius norm of it all

Dividing by the norm without guarding a zero row. An all-zero embedding — an unknown token, a dropped feature — has norm 0, and 0 / 0 gives nan that then spreads through the whole similarity matrix.

Practise it: Normalise Rows to Unit Length

L2 normalisation and cosine similarity

What it is

Cosine similarity is the dot product of unit vectors. Normalise both matrices' rows once, then a single matmul gives the whole (n, m) similarity matrix.

Why it matters

Normalise-then-multiply is both faster and more numerically sensible than dividing an (n, m) product by an outer product of norms, and it is what every vector database does internally.

In the real world: “Customers who liked this also liked…”: comparing product vectors by angle to find similar items, as streaming and shopping sites do.

Visual example

Cosine similarity

Normalise once, multiply once: unit rows turn a cosine similarity into a plain matrix product.

A — raw rows
xy
u34
v02
z00
rows / ||rows||
unit rows
xy
u0.60.8
v01
z00

Row by row

u: ||u|| = 5 → (0.6, 0.8)

v: ||v|| = 2 → (0, 1)

z: ||z|| = 0 → floored, stays 0

Why this matters

Once every row has length 1, a @ b.T is the entire similarity matrix — one BLAS call for millions of pairs.

Common misconception

Wrong idea: the zero row is an edge case you can handle afterwards.

Correct idea: 0 / 0 is nan, and one nan spreads through the whole matrix. Floor the denominator — np.maximum(norm, 1e-12) — and the zero row scores 0 against everything.

  1. 1Measurenorm(A, axis=1, keepdims=True) — one length per row, shape (n, 1).
  2. 2GuardFloor it so an all-zero row divides by a tiny number, not by zero.
  3. 3Multiplya @ b.T gives every pairwise cosine in one call.
def cosine_similarity(A, B):
    a = A / np.maximum(np.linalg.norm(A, axis=1, keepdims=True), 1e-12)
    b = B / np.maximum(np.linalg.norm(B, axis=1, keepdims=True), 1e-12)
    return a @ b.T           # shape (n, m)

Writing the loop version, or the (n, m) division version, and then discovering the zero-vector case. Both problems are fixed once by normalising with a floor on the denominator — a zero row then stays zero and scores 0 against everything, which is the right answer.

Practise it: Cosine Similarity Without Loops, Normalise Rows to Unit Length, Running Share of a Total

Pairwise distances, and the identity that avoids an intermediate

What it is

Squared Euclidean distance expands to ||a||^2 + ||b||^2 - 2 a·b, so a full distance matrix is two row-norm vectors and one matmul — no (n, m, d) intermediate.

Why it matters

The direct broadcast version allocates n * m * d floats, which for two sets of 20,000 vectors in 128 dimensions is hundreds of gigabytes. The identity uses n * m.

In the real world: Finding the nearest delivery hub for every order, or the most similar past fraud case for every new transaction.

def sq_dists(A, B):
    an = (A ** 2).sum(axis=1)[:, None]      # (n, 1)
    bn = (B ** 2).sum(axis=1)[None, :]      # (1, m)
    d2 = an + bn - 2.0 * (A @ B.T)
    return np.maximum(d2, 0.0)              # clip the small negatives away

Taking the square root without clipping first. Cancellation in the subtraction gives tiny negative values for near-identical vectors, and np.sqrt of -1e-16 is nan — a distance matrix whose diagonal is nan.

Practise it: Pairwise Euclidean Distances

solve, inv and einsum

What it is

np.linalg.solve(A, b) solves Ax = b directly; np.linalg.inv builds the inverse; np.linalg.lstsq handles the non-square least-squares case. np.einsum('ij,jk->ik', A, B) names the axes of an arbitrary contraction.

Why it matters

solve is both faster and more accurate than inv(A) @ b, and einsum expresses the contractions — per-example, per-head, weighted — that have no named function.

In the real world: Fitting a simple pricing model that explains sales from price and advertising spend.

x = np.linalg.solve(A, b)                       # not inv(A) @ b
coef, *_ = np.linalg.lstsq(X, y, rcond=None)    # least squares, X not square
np.einsum('ij,ij->i', A, B)                     # a row-wise dot product
np.einsum('bij,bjk->bik', Q, K)                 # a batched matmul, spelled out

Inverting a matrix to solve a system. It is slower and loses precision, and on a near-singular matrix inv returns enormous numbers with no warning while solve at least reports the condition problem.