3. Matrices as functions

Why this layer exists

Every learned layer in a transformer is a matrix multiply, usually followed by something that is not linear. The query, key, value and output projections of attention are matrix multiplies. The two or three projections of the feed-forward block are matrix multiplies with an activation function between them. The final projection from the hidden state back onto the vocabulary is one more. Between them sit a handful of cheap elementwise operations: normalisation, the activation, the softmax inside attention. The learned parameters that are not in the embedding table live almost entirely in weight matrices, and a model’s non-embedding parameter count is, to a close approximation, the sum of their sizes.

Two practical consequences follow, and both come back later in the book. First, reading model code is reading shapes. A transformer implementation is a sequence of arrays with shapes like (batch, seq, d_model) flowing through matrices with shapes like (d_model, d_ff), and most of what goes wrong in such code, at the level of a single line, is a shape that was not what the author thought it was. Second, the arithmetic of inference is dominated by those multiplies. In llama.cpp the operation is ggml_mul_mat, and the backends (CPU, CUDA, Metal, Vulkan) each carry their own hand-tuned kernels for it. When a model generates one token, it reads every weight matrix once and does one multiply-add with each weight, which is why the speed of single-stream decoding tracks memory bandwidth and why quantization (chapter 14) targets the weight matrices: storing each weight in about 4 bits instead of 16 cuts the bytes read per token by more than a factor of three.

Chapter 2 used one matrix–vector product, unit @ unit[q], to score a query against a whole vocabulary. This chapter treats the matrix as the object: a function from vectors to vectors, whose composition is matrix multiplication. The lab builds Linear, a forward-only affine layer whose shape contract, (..., in) to (..., out), is the one the models of Part II use for their projections.

Mechanism

A matrix is a function

A matrix $\mathbf{W}$ of shape (m, k) is a grid of $mk$ numbers, $W_{ij}$ in row $i$ and column $j$. Its product with a vector $\mathbf{x}$ of shape (k,) is a vector $\mathbf{y} = \mathbf{W}\mathbf{x}$ of shape (m,), with entries

$y_i = \sum_{j=1}^{k} W_{ij}\, x_j , \qquad i = 1, \ldots, m .$

Read as a function, $\mathbf{W}$ takes a vector in $\mathbb{R}^k$ and returns a vector in $\mathbb{R}^m$. The number of columns is the input dimension and the number of rows is the output dimension. A (m, k) matrix times a (k,) vector is only defined when the two $k$ agree, and the result has the matrix’s other dimension.

A function $f$ from $\mathbb{R}^k$ to $\mathbb{R}^m$ is linear if, for all vectors $\mathbf{x}, \mathbf{y}$ of shape (k,) and all scalars $\alpha, \beta$,

$f(\alpha\mathbf{x} + \beta\mathbf{y}) = \alpha f(\mathbf{x}) + \beta f(\mathbf{y}) .$

Scale or add the inputs, and the outputs scale and add the same way. Every matrix gives a linear function. Entry $i$ of $\mathbf{W}(\alpha\mathbf{x} + \beta\mathbf{y})$ is

$\sum_{j=1}^{k} W_{ij}\,(\alpha x_j + \beta y_j) = \alpha \sum_{j=1}^{k} W_{ij}\, x_j + \beta \sum_{j=1}^{k} W_{ij}\, y_j ,$

which is entry $i$ of $\alpha\,\mathbf{W}\mathbf{x} + \beta\,\mathbf{W}\mathbf{y}$. The step in the middle is only distributing multiplication over addition and pulling constants out of a sum.

The converse also holds: every linear function from $\mathbb{R}^k$ to $\mathbb{R}^m$ is a matrix. Write $\mathbf{e}_j$ for the vector of shape (k,) with a 1 in position $j$ and zeros elsewhere. Any $\mathbf{x}$ is $x_1\mathbf{e}_1 + \cdots + x_k\mathbf{e}_k$, so by linearity

$f(\mathbf{x}) = x_1 f(\mathbf{e}_1) + \cdots + x_k f(\mathbf{e}_k) .$

A linear function is therefore fixed by what it does to the $k$ vectors $\mathbf{e}_j$. Stack the $k$ outputs $f(\mathbf{e}_j)$, each of shape (m,), side by side as the columns of an (m, k) matrix, and that matrix computes $f$. A linear layer with $k$ inputs and $m$ outputs is, in full generality, $mk$ numbers.

Two readings of the matrix–vector product

The formula for $y_i$ can be read two ways, and both are in constant use.

Row reading. $y_i$ is the dot product of row $i$ of $\mathbf{W}$, shape (k,), with $\mathbf{x}$. Each output coordinate scores the input against one row, which can be read as a learned pattern. This is the reading chapter 2 used: with the unit-length word vectors as rows, the product is the cosine against every word.

Column reading. Write $\mathbf{w}_j$ for column $j$ of $\mathbf{W}$, shape (m,). Then

$\mathbf{W}\mathbf{x} = x_1\mathbf{w}_1 + x_2\mathbf{w}_2 + \cdots + x_k\mathbf{w}_k ,$

a weighted sum of the columns, with the input’s entries as the weights. This is the construction from the converse above, with $\mathbf{w}_j = f(\mathbf{e}_j)$. It says the output always lies in the set of vectors reachable by mixing the columns: a layer cannot produce anything outside that set, whatever its input.

The two readings compute the same numbers because they are the same double sum taken in a different order. Entry $i$ of the column reading is $\sum_j x_j (\mathbf{w}_j)_i = \sum_j x_j W_{ij}$, which is the row reading’s $y_i$. A small case, with $\mathbf{W}$ of shape (2, 3) and $\mathbf{x} = (1, 1, 2)$:

$\mathbf{W} = \begin{pmatrix} 1 & 2 & 0 \\ 0 & 1 & -1 \end{pmatrix}, \qquad \mathbf{W}\mathbf{x} = \begin{pmatrix} 1 + 2 + 0 \\ 0 + 1 - 2 \end{pmatrix} = 1\begin{pmatrix} 1 \\ 0 \end{pmatrix} + 1\begin{pmatrix} 2 \\ 1 \end{pmatrix} + 2\begin{pmatrix} 0 \\ -1 \end{pmatrix} = \begin{pmatrix} 3 \\ -1 \end{pmatrix} .$

The middle expression is two dot products, one per row; the right-hand one is three scaled columns.

Matrix multiplication is composition

Take two linear functions in sequence: $\mathbf{B}$ of shape (k, m) maps $\mathbb{R}^m$ to $\mathbb{R}^k$, and $\mathbf{A}$ of shape (n, k) maps $\mathbb{R}^k$ to $\mathbb{R}^n$. For an input $\mathbf{x}$ of shape (m,), the composite is $\mathbf{A}(\mathbf{B}\mathbf{x})$, shape (n,). Write it out, calling the intermediate $\mathbf{z} = \mathbf{B}\mathbf{x}$, shape (k,):

$\big(\mathbf{A}(\mathbf{B}\mathbf{x})\big)_i = \sum_{p=1}^{k} A_{ip}\, z_p = \sum_{p=1}^{k} A_{ip} \sum_{j=1}^{m} B_{pj}\, x_j = \sum_{j=1}^{m} \Big( \sum_{p=1}^{k} A_{ip}\, B_{pj} \Big) x_j .$

The last step swaps the order of two finite sums. The bracket depends on $i$ and $j$ but not on $\mathbf{x}$, so it is the entry in row $i$, column $j$ of a single (n, m) matrix that does the whole job. That matrix is the product $\mathbf{A}\mathbf{B}$:

$(\mathbf{A}\mathbf{B})_{ij} = \sum_{p=1}^{k} A_{ip}\, B_{pj} ,$

shape (n, k) times (k, m) gives (n, m). The definition is not arbitrary: it is exactly what makes $(\mathbf{A}\mathbf{B})\mathbf{x} = \mathbf{A}(\mathbf{B}\mathbf{x})$ hold. Entry $(i, j)$ is the dot product of row $i$ of $\mathbf{A}$ with column $j$ of $\mathbf{B}$, both of length $k$, which is why the inner dimensions must agree.

Three facts follow. The product is associative, $(\mathbf{A}\mathbf{B})\mathbf{C} = \mathbf{A}(\mathbf{B}\mathbf{C})$, because both sides are the composite function “apply $\mathbf{C}$, then $\mathbf{B}$, then $\mathbf{A}$“, and a linear function has only one matrix. It is not commutative: $\mathbf{A}\mathbf{B}$ and $\mathbf{B}\mathbf{A}$ are different functions in general, and are usually not both defined. And a vector is the special case $m = 1$: $\mathbf{x}$ of shape (k,) behaves as a (k, 1) matrix, and the matrix–vector product is the matrix–matrix product with one column.

Depth without a nonlinearity buys nothing

Stack two linear layers, $\mathbf{W}_1$ of shape (h, k) and $\mathbf{W}_2$ of shape (m, h), on an input $\mathbf{x}$ of shape (k,). By the composition result,

$\mathbf{W}_2(\mathbf{W}_1\mathbf{x}) = (\mathbf{W}_2\mathbf{W}_1)\,\mathbf{x} ,$

and $\mathbf{W}_2\mathbf{W}_1$ is one matrix of shape (m, k). Two layers compute a function that one layer could have computed. By induction the same holds for any depth: a stack of linear layers is one linear layer. It can be worse than one layer, too. Every output of the stack is $\mathbf{W}_2$ applied to something, so by the column reading it is a mix of the $h$ columns of $\mathbf{W}_2$, each of shape (m,). If $h < k$ and $h < m$, the stack can therefore reach at most an $h$-dimensional slice of $\mathbb{R}^m$, a restriction a single (m, k) layer does not have.

This is why every block in a transformer puts something non-linear between its matrices. The feed-forward block is “project up, apply an elementwise non-linear function, project down”; without the middle step it would collapse to a single (d, d) matrix. Chapter 9 introduces the non-linearity and builds the first multi-layer model on it.

Affine maps and the bias

A linear function sends zero to zero: $f(\mathbf{0}) = f(0 \cdot \mathbf{0}) = 0 \cdot f(\mathbf{0}) = \mathbf{0}$. A layer that must output something other than zero for a zero input needs a constant added:

$f(\mathbf{x}) = \mathbf{W}\mathbf{x} + \mathbf{b} ,$

with $\mathbf{W}$ of shape (m, k), $\mathbf{x}$ of shape (k,), and the bias $\mathbf{b}$ of shape (m,). A function of this form is affine. It is not linear when $\mathbf{b} \neq \mathbf{0}$, since $f(\mathbf{x} + \mathbf{y}) = \mathbf{W}\mathbf{x} + \mathbf{W}\mathbf{y} + \mathbf{b}$ has one bias where $f(\mathbf{x}) + f(\mathbf{y})$ has two. Subtracting $f(\mathbf{0}) = \mathbf{b}$ recovers the linear part, $f(\mathbf{x}) - f(\mathbf{0}) = \mathbf{W}\mathbf{x}$, and that is exactly how the lab test test_linear_is_affine checks a layer. The composition argument carries over: two affine maps in a row are one affine map (exercise e), so biases do not rescue depth either. The Llama family of models drops the biases from its projections altogether; the layer in this chapter keeps them, since they cost $m$ parameters against $mk$ for the weights.

Transpose

The transpose of $\mathbf{A}$, shape (n, k), is $\mathbf{A}^\top$, shape (k, n), with rows and columns exchanged: $(\mathbf{A}^\top)_{ij} = A_{ji}$. Treating a vector $\mathbf{a}$ of shape (d,) as a (d, 1) column, $\mathbf{a}^\top$ is a (1, d) row, and the product $\mathbf{a}^\top\mathbf{b}$ is a (1, 1) matrix whose single entry is $\sum_i a_i b_i$: the dot product of chapter 2, written as a matrix product.

Transposes appear in model code whenever two conventions meet. The rule that governs them, $(\mathbf{A}\mathbf{B})^\top = \mathbf{B}^\top\mathbf{A}^\top$, is exercise (b).

Batches and the row-vector convention

The mathematics above writes vectors as columns and puts the matrix on the left: $\mathbf{W}\mathbf{x}$ with $\mathbf{W}$ of shape (out, in). Code almost always does the opposite. An input is a row, the weight is stored as (in, out), and the layer computes

y = x @ W + b    # (in,) @ (in, out) + (out,) -> (out,)

The reason is batching. Stack $B$ input rows into a matrix $\mathbf{X}$ of shape (B, in). Then $\mathbf{X}\mathbf{W}$ has shape (B, out), and the product formula gives its entry in row $r$, column $j$ as

$(\mathbf{X}\mathbf{W})_{rj} = \sum_{p=1}^{\text{in}} X_{rp}\, W_{pj} ,$

which involves only row $r$ of $\mathbf{X}$. So row $r$ of $\mathbf{X}\mathbf{W}$ is row $r$ of $\mathbf{X}$, shape (in,), times $\mathbf{W}$, exactly as if that input had been processed alone. One matrix multiply applies the layer to every row independently, and a batch of one is the same code as a batch of a thousand. Putting the batch on the first axis also keeps each input row contiguous in memory.

The two conventions compute the same thing. The row form $\mathbf{x}^\top\mathbf{W}$, with $\mathbf{W}$ of shape (in, out), has the same entries as the column form $\mathbf{W}^\top\mathbf{x}$ (exercise b), so a weight stored in one convention is the transpose of the same weight stored in the other. PyTorch’s nn.Linear stores its weight as (out_features, in_features) and computes x @ W.T + b; the Linear in this book stores (in_features, out_features) and computes x @ W + b. Both are the same layer, and a checkpoint converted between the two conventions has its weight matrices transposed.

Broadcasting

In x @ W + b, the left operand of + has shape (B, out) and the right has shape (out,). The addition works because of broadcasting, numpy‘s rule for elementwise operations on arrays of different shapes. Stated precisely:

1. Line the two shapes up at their trailing (rightmost) axes. If one shape has fewer axes, pad it on the left with axes of size 1.
2. Compare each aligned pair of sizes. They are compatible if they are equal or if one of them is 1. Any other pair is an error.
3. The result has, on each axis, the larger of the two sizes. An operand whose size on an axis is 1 is repeated along that axis. No copy is made; the repeat is done with a stride of zero.

For the bias: (B, out) and (out,) align as (B, out) and (1, out); the pairs are $(B, 1)$, compatible, and $(\text{out}, \text{out})$, equal; the result is (B, out), and the same bias row is added to every row of the batch.

The rules also explain the silent failure chapter 2 warned about. (3,) + (1,) is legal, because a size-1 axis broadcasts against anything, and returns shape (3,). A worse case is a column meeting a row: a vector of per-example losses with shape (B,) added to targets of shape (B, 1) aligns as (1, B) against (B, 1), and both pairs contain a 1, so the result is a (B, B) matrix of every loss plus every target. No error is raised, and a following .mean() returns a plausible-looking number. The defence is to check shapes explicitly at the boundaries where they are known, which is what assert_shape is for.

Leading axes

Language-model activations have more than two axes: a batch of $B$ sequences of $T$ tokens, each a vector of width $d$, is an array of shape (B, T, d). numpy‘s @ (and np.matmul) handles this by treating everything but the last two axes as a stack of matrices, and broadcasting across the stack with the rules above. For the linear layer’s case, (B, T, in) @ (in, out), the two-axis weight broadcasts against every one of the $B$ matrices of shape (T, in) and the result has shape (B, T, out). A one-axis left operand is promoted to a one-row matrix for the product and the added axis is removed afterwards, so (in,) @ (in, out) gives (out,).

The consequence is that a linear layer acts on the last axis only. Whatever leading axes the input has, batch, sequence position, attention head, pass through unchanged. The layer is written once, for shape (..., in) to (..., out), and the same code serves an input of any rank. The projections in the models of Part II follow this contract.

Initialization scale

A layer’s weights have to start somewhere, and the scale of the starting values matters. Each output coordinate is a sum of in products $x_j W_{ji}$. If the weights were of size about 1, that sum would grow with the number of terms, and a stack of layers would multiply the growth layer by layer until the activations overflowed or, with weights too small, shrank to zero. The remedy used here is to draw each weight and bias uniformly from $[-1/\sqrt{\text{in}}, 1/\sqrt{\text{in}}]$. With that scale, an input whose entries have size about 1 produces outputs whose entries have size about 1 (somewhat smaller, by a constant factor), whatever the input width. It is also PyTorch’s default for nn.Linear. Why $\sqrt{\text{in}}$ and not $\text{in}$ is a statement about the variance of a sum of independent random terms; chapter 9 derives it once probability (chapter 4) supplies the tools.

Walkthrough

The reference module is py/tinygpt/linear.py.

def assert_shape(x, *dims):
    """Raise ValueError unless x has exactly len(dims) axes matching dims; None matches any size."""
    shape = np.shape(x)
    if len(shape) != len(dims) or any(d is not None and d != s for d, s in zip(dims, shape)):
        raise ValueError(f"expected shape {dims}, got {shape}")

assert_shape is a shape check for function boundaries. The expected shape is given as separate arguments, with None for an axis whose size varies, so assert_shape(x, None, 3) means “a matrix with three columns and any number of rows”. np.shape works on lists as well as arrays, so the check needs no conversion.

The rank comparison comes first and is not redundant. zip stops at the shorter of its two arguments, so without the len(shape) != len(dims) test a (2, 3) array would pass a check for (2,) and a check for (2, 3, 1): the zip would compare the axes the two shapes share and ignore the rest. Both are among the rejections the lab tests. The function raises ValueError, the exception numpy itself uses for incompatible shapes, and the message prints both shapes, which is the information needed to find the bug.

def matmul(a, b):
    """(n, k) @ (k, m) -> (n, m): out[i, j] is the dot product of row i of a and column j of b."""
    a, b = np.asarray(a, dtype=np.float64), np.asarray(b, dtype=np.float64)
    if a.ndim != 2 or b.ndim != 2:
        raise ValueError(f"matmul needs two matrices, got {a.shape} and {b.shape}")
    n, k = a.shape
    if b.shape[0] != k:
        raise ValueError(f"inner dimensions differ: {a.shape} @ {b.shape}")
    m = b.shape[1]
    out = np.zeros((n, m))
    for i in range(n):
        for j in range(m):
            for p in range(k):
                out[i, j] += a[i, p] * b[p, j]
    return out

matmul is the index formula written as three loops: $i$ over rows of the output, $j$ over its columns, $p$ over the shared inner dimension. a and b are converted to float64 for the same reason dot did it in chapter 2. The shape checks reject anything that is not a pair of matrices with matching inner dimensions. A vector is rejected rather than promoted, which is stricter than @: np.zeros(3) @ np.zeros((3, 2)) returns shape (2,), while matmul raises. The strictness is the point of a from-scratch version whose job is to make every shape visible.

The innermost statement runs once per $(i, j, p)$ triple, so matmul does exactly $nkm$ multiply-adds. Cost is conventionally counted in FLOPs, floating-point operations, with a multiply-add counting as two: $2nkm$ FLOPs for one product.

Apply that to one token passing through a square layer of width $d$: the input has shape (1, d), the weight (d, d), so $n = 1$, $k = m = d$, and the product costs $d^2$ multiply-adds, $2d^2$ FLOPs. For $d = 4096$ that is 16,777,216 multiply-adds, about $3.4 \times 10^7$ FLOPs, for one token through one projection. The pattern generalises. Every weight in the layer takes part in exactly one multiply-add per token, so the FLOPs for one token through any weight matrix are twice its number of weights, and the FLOPs for one token through a whole model are, to first order, twice its non-embedding parameter count. The embedding table is excluded because looking up a token’s row copies it and does no multiply-adds. The biases, normalisations and activations add work proportional to the widths, not their products; the attention scores add work that grows with context length (chapter 10). For a batch of $n$ tokens the matrix multiplies cost $n$ times as much, but the weights are still read from memory once per batch, not once per token. That ratio, arithmetic per byte of weight read, is what chapter 17’s batching improves.

class Linear:
    """y = x W + b, applied to the last axis of x: (..., in) -> (..., out)."""

    def __init__(self, in_features, out_features, rng):
        bound = 1.0 / np.sqrt(in_features)
        self.weight = rng.uniform(-bound, bound, size=(in_features, out_features))
        self.bias = rng.uniform(-bound, bound, size=(out_features,))

    def __call__(self, x):
        x = np.asarray(x, dtype=np.float64)
        if x.shape[-1] != self.weight.shape[0]:
            raise ValueError(f"expected last axis {self.weight.shape[0]}, got shape {x.shape}")
        return x @ self.weight + self.bias

Linear is the affine map of the mechanism in the row-vector convention: the weight is stored (in_features, out_features) so the forward pass is x @ W + b with no transpose. Both parameters are drawn from rng.uniform(-bound, bound) with bound = 1/sqrt(in_features). The caller passes the random generator in, so two layers built from generators with the same seed are identical, and tests and training runs are reproducible without any global state.

__call__ accepts any input whose last axis has width in_features. It checks x.shape[-1] directly rather than calling assert_shape, because assert_shape fixes the number of axes and a layer should accept (in,), (B, in) and (B, T, in) alike. The product then follows the rules for leading axes, (..., in) @ (in, out) giving (..., out), and the bias, shape (out,), broadcasts across every leading axis. Without the explicit check, a wrong-width input would still fail inside @, but with a message about matmul core dimensions rather than one that names the layer’s expected width.

Exercises

(a) What is the shape of each intermediate result, and of the final result, when evaluating (8, 16) @ (16, 32) @ (32, 4) left to right? How many multiply-adds does it take, and how many if evaluated right to left?

(b) Let $\mathbf{x}$ have shape (in,), treated as a (in, 1) column, and $\mathbf{W}$ shape (in, out). Show that the row-vector product $\mathbf{x}^\top\mathbf{W}$, which is what x @ W computes, equals $(\mathbf{W}^\top\mathbf{x})^\top$. Prove the general rule $(\mathbf{A}\mathbf{B})^\top = \mathbf{B}^\top\mathbf{A}^\top$ on the way.

(c) Which of these broadcasts succeed, and what is the result shape of each: (3, 1) + (1, 4), (3,) + (4,), (2, 3, 4) + (3, 1)?

(d) A feed-forward projection maps width 4096 to width 11008. Count the FLOPs of one forward pass through it, including the bias, for 1 token and for a batch of 512 tokens.

(e) Prove that the composition of two affine maps is affine, and write the combined weight and bias. Use the row-vector convention of the code: $f_1(\mathbf{x}) = \mathbf{x}\mathbf{W}_1 + \mathbf{b}_1$ and $f_2(\mathbf{h}) = \mathbf{h}\mathbf{W}_2 + \mathbf{b}_2$, with $\mathbf{x}$ of shape (1, k), $\mathbf{W}_1$ (k, h), $\mathbf{b}_1$ (1, h), $\mathbf{W}_2$ (h, m), $\mathbf{b}_2$ (1, m).

Lab

Implement assert_shape, matmul and Linear in py/labs/ch03/starter.py, then run

make lab CH=03 IMPL=mine

until it passes. Read py/labs/ch03/test_lab.py first. It pins the details the prose leaves to judgement: assert_shape must reject a rank mismatch in either direction, not only a size mismatch; matmul must raise ValueError for a vector argument and for mismatched inner dimensions, and must match @ on random input; Linear must accept inputs with any number of leading axes, reject a wrong last-axis width with ValueError, and draw both weight and bias within $\pm 1/\sqrt{\text{in}}$. test_linear_is_affine subtracts the layer’s output at zero before checking additivity, since with a bias the layer itself is not linear.

Then measure what the explicit loops cost. Save this script outside the repository, for example as /tmp/matmul_timing.py:

import time

import numpy as np

from tinygpt.linear import matmul


def seconds(f, repeat):
    """Best wall-clock time of `repeat` calls to f."""
    best = float("inf")
    for _ in range(repeat):
        t0 = time.perf_counter()
        f()
        best = min(best, time.perf_counter() - t0)
    return best


rng = np.random.default_rng(0)
for n in [64, 256]:
    a, b = rng.normal(size=(n, n)), rng.normal(size=(n, n))
    np.testing.assert_allclose(matmul(a, b), a @ b)
    loops = seconds(lambda: matmul(a, b), repeat=1)
    blas = seconds(lambda: a @ b, repeat=5)
    print(f"n={n:3d}  matmul {loops:8.3f} s   @ {blas * 1e6:8.1f} us   ratio {loops / blas:>9,.0f}")

and run it from py/ with PYTHONPATH=. uv run python /tmp/matmul_timing.py, changing the import to from labs.ch03.starter import matmul to time yours. @ is timed as the best of five calls because its first call can include one-off start-up costs; the loop version is timed once, since a single 256×256 run already takes seconds. One run on one machine, a laptop with an Intel Core i3-1315U, printed:

n= 64  matmul    0.273 s   @     67.9 us   ratio     4,020
n=256  matmul   22.331 s   @   1462.5 us   ratio    15,269

These figures are one run, not a benchmark. The machine was busy with other work at the time, and two further runs of the same script gave ratios between about 1,000 and about 20,000, mostly because the time for @ moved by a factor of ten. Your hardware will give different numbers again. The shape of the result is what to explain.

• The loop version scales as $n^3$. Going from 64 to 256 multiplies $n$ by 4 and the work by $4^3 = 64$: from 262,144 to 16,777,216 multiply-adds. The measured time grew by a factor of 82, close to 64 given the noise. Dividing time by multiply-adds gives the cost of one trip through the inner loop body: about 1.0 microseconds at $n = 64$ and 1.3 at $n = 256$.
• Each trip through that body is interpreted. out[i, j] += a[i, p] * b[p, j] runs as a sequence of bytecode instructions. Each of the three reads, out[i, j], a[i, p] and b[p, j], parses an index tuple and allocates a new numpy scalar object for the value; then come a multiply, an add, and a store back into the array, each dispatched at run time on the types of its operands. The one floating-point multiply-add the line exists for is a tiny fraction of that work.
• @ calls BLAS. numpy hands the product to a compiled BLAS library (on Linux x86-64 the numpy wheels ship OpenBLAS, which is what the run above used; the macOS arm64 wheels use Apple’s Accelerate). Its kernels are machine code with no per-element dispatch, use SIMD instructions that do several multiply-adds at once, split the matrices into blocks sized to fit in cache so each loaded number is reused many times, and can spread the work across cores. At 256×256 the run above works out to $2 \cdot 256^3$ FLOPs in 1.46 ms, about 23 billion floating-point operations per second.

The ratio, three to four orders of magnitude, is the reason no later chapter calls matmul: it exists to make the index formula concrete, and @ is what the models use. It is also the reason llama.cpp spends its engineering effort on matrix-multiply kernels for each backend: the arithmetic is fixed by the model, and how fast it runs depends on how the kernel uses the hardware.

Further reading

• Gilbert Strang, Introduction to Linear Algebra (5th ed., 2016), chapters 1–2. Matrices acting on vectors, the row and column pictures of a matrix–vector product, and the rules of matrix multiplication.
• NumPy documentation, “Broadcasting”. The authoritative statement of the rules above, with worked shape examples.
• Supplementary video series: 3Blue1Brown, Essence of Linear Algebra (2016), the chapters on linear transformations and on matrix multiplication as composition. The function view of a matrix, drawn in motion.