Math · Week 1 · Note 3 of 14

Matrices, matrix multiplication and shapes

A matrix is a table of numbers, and multiplying one by a vector is doing a dot product with every row at once. That single idea, plus a rule about which shapes fit together, is nearly all the matrix algebra machine learning uses. This note builds it with small tables and ends with the shape bookkeeping that catches most bugs.

Why this note matters

Your data is a table: one row per house, one column per feature. Your model’s weights are a list, one per feature. “Predict every house at once” is the table times the list, and that is a matrix-vector product. The linear regression note already used this, with X @ w; this note is where you learn exactly what that line does and when it fails. Later, every piece of a neural network ? is a matrix multiplication followed by something simple, the number everyone quotes for a model, its parameter count, is the sizes of its matrices added up, and the error you will see most often while learning, some version of “shapes (32, 3) and (4, 3) not aligned”, is a matrix whose rows and columns do not fit the thing next to it.

Interview questions that are really this note:

• “What is the shape of the result when a (32, 3) data matrix meets a (4, 3) weight matrix, and why does it fail as written?”
• “A weight matrix from 768 inputs to 3072 outputs: what shape is it, and how many parameters is that?”
• “How many arithmetic operations does one matrix multiply cost? Why do GPUs help?”
• “Why is $AB$ different from $BA$, and when does the bracketing $(AB)C$ versus $A(BC)$ matter?”
• “Why is attention ? quadratic in the length of the text?” (Week 8; the shape half of the answer is here.)

In plain words

What a matrix is

A matrix is a rectangular grid of numbers with $m$ rows and $n$ columns; we say its shape is $m \times n$, read “m by n”, rows first. Matrices get capital letters, $A$ or $W$, and entries get two subscripts, $A_{ij}$ for row $i$ and column $j$. Math counts from 1; NumPy and PyTorch count from 0, so the 2 in Fig. 1 is $A_{12}$ on paper and A[0, 1] in code.

A vector fits into this picture as a matrix with one column: a column vector of shape $n \times 1$. Written flat, it is a row vector of shape $1 \times n$. Mathematicians default to columns, which is why formulas say $W\mathbf{x}$ with the vector on the right. NumPy adds a third option, the plain 1-D array of shape (n,) with no second number at all, and treats it as whichever one makes the multiplication work. Most shape bugs come from mixing up these three; the common mistakes list has the specifics.

You will meet matrices in three roles, and it helps to name the role each time:

• A data matrix $X$: one row per example, one column per feature. Four days of walk inputs is a $4 \times 3$ matrix. A batch of 32 images flattened to 784 pixels is $32 \times 784$.
• A weight matrix $W$: one row per output the model computes, one column per input feature. The two walk scores from the vectors note make a $2 \times 3$ weight matrix.
• A lookup table: a language model keeps one embedding ? vector per token ? it knows. For GPT-2 that is a $50{,}257 \times 768$ matrix, and “looking up a token” means reading its row.

Matrix times vector: a stack of dot products

The matrix-vector product $W\mathbf{x}$ is defined entry by entry: entry $i$ of the result is the dot product of row $i$ of $W$ with $\mathbf{x}$.

Two consequences fall straight out of the definition. First, the vector must have as many entries as each row has, so a $2 \times 3$ matrix can only multiply a 3-vector. Second, the result has one entry per row, so it is a 2-vector. In shape language: $(2 \times 3)(3 \times 1) = (2 \times 1)$. The 3s touch and cancel; the 2 and the 1 survive. That is the shape rule, and the next section says it in full.

There is a second way to read the same product, and it is worth seeing once because it is how 3Blue1Brown and many textbooks present it. Instead of “each row dotted with $\mathbf{x}$”, read it as “the columns of $W$, mixed in the amounts given by $\mathbf{x}$”:

The result is a linear combination of the columns, with the input’s entries as the coefficients. Both readings give the same numbers; use the row reading for computing and the column reading for the “matrix as a function” picture in a moment.

Matrix times matrix and the shape rule

For $A$ of shape $m \times n$ and $B$ of shape $n \times p$, the product $C = AB$ has shape $m \times p$, and each entry is a dot product of a row of $A$ with a column of $B$:

Work one fully. Say every step out loud the first time; it is the only way the row-column dance becomes automatic.

The habit that prevents bugs: write the shapes before you compute. “(2 × 3) times (3 × 2): inner 3s match, result is 2 × 2.” Do this for every multiplication in every formula until it is reflex. Interviewers watch for it, and it is how you will debug your own code: print the shapes, line them up, find the pair that does not touch.

Order matters, grouping does not

Not commutative. $AB \ne BA$ in general. With the $A$ and $B$ above, $AB$ is $2 \times 2$ but $BA$ is $3 \times 3$: not even the same shape. Even for square matrices the entries differ:

(Multiplying by that 0-1 matrix on the right swaps the columns of the other matrix; on the left it swaps the rows. Which side you multiply on decides which axis you act on. This is why formulas are careful to write $XW^\top$ and not $W^\top X$.)

Associative. $(AB)C = A(BC)$, always. The answer is identical; the cost is not. Suppose a big $768 \times 768$ matrix is written as the product of two thin ones, $B$ ($768 \times 8$) and $A$ ($8 \times 768$), so the big matrix is $BA$. (Week 10 meets a technique called LoRA that does exactly this to shrink the cost of adapting a language model.) To apply it to one input vector $\mathbf{x}$ you can compute:

• $(BA)\mathbf{x}$: first form $BA$, which is $768 \cdot 8 \cdot 768 \approx 4.7$ million multiplications, then multiply the vector.
• $B(A\mathbf{x})$: first $A\mathbf{x}$ ($8 \times 768 = 6{,}144$ multiplications, giving an 8-vector), then $B$ times that (another 6,144). About 12 thousand in total.

Same result, four hundred times less work. Choosing the bracketing is a real engineering decision, and associativity is what makes it legal.

Distributive. $A(B + C) = AB + AC$. This is what lets you write $(W + BA)\mathbf{x} = W\mathbf{x} + BA\mathbf{x}$ and compute the two halves separately, which is how that week-10 technique adds its small update to a frozen big matrix.

The identity. The square matrix with 1s on the diagonal and 0s elsewhere, $I$, leaves anything it multiplies unchanged: $I\mathbf{x} = \mathbf{x}$ and $IA = AI = A$. It is the matrix version of the number 1. It shows up in week 3, when a penalty on big weights adds a multiple of $I$ to a matrix, and again in several deep-learning tricks later.

A matrix is a function

A matrix defines a function $\mathbf{x} \mapsto M\mathbf{x}$, and it is a special kind of function called a linear transformation: it respects addition and scaling, $M(\mathbf{u} + \mathbf{v}) = M\mathbf{u} + M\mathbf{v}$ and $M(c\mathbf{v}) = cM\mathbf{v}$. One consequence matters a great deal in week 5: composing two linear functions gives a linear function, and in matrix language composing the functions is multiplying the matrices. Apply $A$ then $B$ and you have applied $BA$ (note the order: the first applied is on the right, nearest the vector). So two matrix multiplications in a row, with nothing in between, are no more powerful than one. That is why neural networks put a simple bend between their layers.

Read off where the two basis arrows go and you know the whole matrix. In Fig. 4, “one step right” lands at $(2, 0)$ and “one step up” lands at $(1, 1)$; those are the two columns of $M$. Any other point is a combination of those two steps, so it goes to the same combination of the two landing spots: $(1, 1)$ is one right plus one up, so it goes to $(2, 0) + (1, 1) = (3, 1)$. This is the column reading from the matrix-vector section, now as a picture.

The picture is what lets you think about a weight matrix as something that reshapes the space the data lives in, so that the categories become easier to separate. The eigenvector note in week 4 asks the natural follow-up: are there directions a given matrix only stretches, without turning? That question is behind PCA ?.

Batches: many inputs at once

Four walk days as a data matrix, the two scores as a weight matrix, the biases as a vector:

We want each day (a row of $X$) dotted with each score’s weights (a row of $W$). Row-with-row is not what matmul does; it does row-with-column. So flip $W$ on its side first, turning its rows into columns, and multiply:

Check the shapes: $(4 \times 3)(3 \times 2) = 4 \times 2$. Check one entry: day 4 is $(1, 0, 1)$, the walk weights are $(3, -2, -4)$, dot product $-1$, plus bias $-1$, gives $-2$, the bottom-left entry. Check the bias: it is a 2-vector, one per score, and it was added to every row. That silent stretching of a $(2,)$ vector across a $(4, 2)$ matrix is broadcasting, covered below.

This “weights times a batch, plus a bias” step is so common that deep-learning libraries give it a name, a linear layer, and store $W$ with shape (outputs, inputs), one row per output, computing $XW^\top + \mathbf{b}$ exactly as above. (You will meet it as nn.Linear in week 6.) Two facts to memorize now, because they answer a whole family of interview questions:

• The weight shape is (outputs, inputs). A layer from 768 inputs to 3072 outputs has a weight of shape (3072, 768). Parameters: $3072 \times 768 + 3072 = 2{,}362{,}368$, the weights plus one bias per output.
• It changes only the last axis. Input (batch, in) becomes (batch, out). Input (batch, seq, in), which is how a language model carries a batch of sentences, becomes (batch, seq, out). Every leading axis is left alone.

Two of these steps in a row, with a bend between them, as a shape table. Nothing in it is new; it is the shape rule applied twice. (In week 5 this table is a small neural network.)

Transpose, and all pairwise dot products

The transpose $A^\top$ of an $m \times n$ matrix is the $n \times m$ matrix with $(A^\top)_{ij} = A_{ji}$:

Three rules, each checkable by shapes alone:

• $(A^\top)^\top = A$. Flip twice, back where you started.
• $(AB)^\top = B^\top A^\top$. The order reverses. Check: $A$ is $2 \times 3$, $B$ is $3 \times 2$, so $(AB)^\top$ is $2 \times 2$; and $B^\top A^\top$ is $(2 \times 3)(3 \times 2) = 2 \times 2$. The other order, $A^\top B^\top$, would be $(3 \times 2)(2 \times 3)$, a different shape entirely.
• $\mathbf{x}^\top \mathbf{y}$ is the dot product. A column $\mathbf{x}$ ($n \times 1$) transposed is a row ($1 \times n$), times a column ($n \times 1$) gives $1 \times 1$: one number. This is why papers write $\mathbf{w}^\top\mathbf{x}$ for $\mathbf{w}\cdot\mathbf{x}$.

Fig. 5 is the single most important shape in the later tracks, so make sure it is clear now. Put $n$ vectors of dimension $d$ as the rows of a matrix $A$ ($n \times d$) and another $n$ as the rows of $B$ (also $n \times d$). Then $AB^\top$ is $(n \times d)(d \times n) = n \times n$, and by the definition of matmul its entry $(i, j)$ is row $i$ of $A$ dotted with column $j$ of $B^\top$, which is row $j$ of $B$. So entry $(i, j)$ is $\mathbf{a}_i \cdot \mathbf{b}_j$: the similarity between vector $i$ of the first stack and vector $j$ of the second. All $n^2$ pairs in one multiplication.

The same trick with one matrix and itself, $XX^\top$, gives every example’s dot product with every other example, a Gram matrix. It is how you would compute all pairwise similarities between a thousand embeddings ? in one line: E @ E.T. And in week 8, the step inside a language model that decides how much each word should look at each other word is exactly Fig. 5, with one stack called the queries and the other the keys.

How much work a multiplication is

Multiplying $m \times n$ by $n \times p$ computes $mp$ entries, each a dot product of length $n$: $mnp$ multiplications and about $mnp$ additions, so roughly $2mnp$ FLOPs (floating-point operations). For one of GPT-2’s bigger weight matrices, 768 by 3072, applied to a text of 1,024 tokens ?, that is $2 \times 1024 \times 768 \times 3072 \approx 4.8$ billion FLOPs for one matrix of one pass through the model. Multiply by the number of such matrices, the number of training tokens and a factor of three for the learning step, and you have the training cost estimates you see in papers.

Three facts that follow:

• Why GPUs. The $mp$ dot products do not depend on each other, so they can all run at the same time. A GPU is thousands of small arithmetic units built to do exactly that, and modern ones have dedicated matrix-multiply hardware (tensor cores).
• Why memory, not arithmetic, is often the limit. Reading $mn + np$ numbers and writing $mp$ can take longer than the arithmetic when $n$ is small. Batching exists largely so that each weight read from memory is used for many rows at once.
• Why parameter count ≈ FLOPs per token. A weight matrix of $mn$ entries costs $2mn$ FLOPs per input vector. So running a model with $N$ parameters once costs about $2N$ FLOPs per token. That rule of thumb answers a whole family of interview questions about the cost of using a model.

Element-wise operations and broadcasting

Element-wise operations pair up matching entries and never mix rows with columns. $A + B$, $A \odot B$ (A * B in code), $\max(0, A)$, $A / 2$ are all element-wise, and all keep the shape. The trap: A * B in NumPy is element-wise, not matrix multiplication. Matmul is A @ B. In the last note the same trap was w * x versus w @ x; here it costs you a wrong answer with no error message, because two matrices of the same shape can be either *-ed or, if square, @-ed.

Broadcasting is the rule for element-wise operations between arrays of different shapes. Line the shapes up from the right. Two dimensions are compatible if they are equal or one of them is 1 (or missing). The size-1 dimension is stretched to match. So:

The last row is the one that bites. To subtract a per-row mean of shape (4,) from a (4, 2) matrix you must first make it (4, 1) with means[:, None] or keepdims=True. Standardising features in the feature scaling note does exactly this kind of thing, and so does a normalisation step inside every language model.

Code

Type these out. Before running each line that prints a shape, write down the shape you expect.

Shapes, indexing, transpose and the error you will see

import numpy as np

A = np.array([[1, 2, 3],
              [4, 5, 6]])      # 2 rows, 3 columns
B = np.array([[7, 8],
              [9, 10],
              [11, 12]])       # 3 rows, 2 columns

print(A.shape)                 # (2, 3)
print(B.shape)                 # (3, 2)

print(A @ B)                   # @ is matrix multiplication
# [[ 58  64]
#  [139 154]]                  shape (2, 2): the inner 3s matched

# the other order gives a different, 3 x 3 answer
print(B @ A)
# [[ 39  54  69]
#  [ 49  68  87]
#  [ 59  82 105]]

# transpose: rows become columns, shape (3, 2)
print(A.T)
# [[1 4]
#  [2 5]
#  [3 6]]

print(A[0])                    # row 0:     [1 2 3]
# column 1:  [2 5]   (the colon means "every row")
print(A[:, 1])
print(A[1, 2])                 # row 1, column 2, counting from 0:  6

# Shapes that do not fit raise an error. try / except catches it so we can
# print the message.
try:
    # (2, 3) @ (2, 3): the inner sizes 3 and 2 do not match
    A @ A
except ValueError as err:
    print(err)
# matmul: Input operand 1 has a mismatch in its core dimension 0, ... (size
# 2 is different from 3)

A batch through a weight matrix: the walk days and the houses

import numpy as np

# Four days (rows) x three features (columns): sunny, cold, tired.
X = np.array([[1, 0, 0],
              [1, 1, 0],
              [0, 0, 1],
              [1, 0, 1]], dtype=float)     # shape (4, 3)

# Two scores, one row of weights each: shape (2, 3) = (outputs, inputs).
W = np.array([[3, -2, -4],                 # score 1: "should I walk?"
              [1, -1,  0]], dtype=float)   # score 2: "is the weather good?"
b = np.array([-1, 0], dtype=float)         # one bias per score: shape (2,)

# (4, 3) @ (3, 2) -> (4, 2); b is added to every row
Z = X @ W.T + b
print(Z)
# [[ 2.  1.]
#  [ 0.  0.]
#  [-5.  0.]
#  [-2.  1.]]
# : one row per day, one column per score
print(Z.shape)
# prints (4, 2)

# The house data the same way: three houses, two features (size, rooms).
houses = np.array([[1.0, 2],
                   [1.5, 3],
                   [2.0, 3]])              # shape (3, 2)
w = np.array([150.0, 15.0])                # one weight per feature
b0 = 10.0                                  # one bias
# @ (2,) -> (3,): all three predictions at once
print(houses @ w + b0)
# prints (3, 2)
# [190. 280. 355.]

The shape table of two steps with a bend between

import numpy as np

# a random-number generator; the 0 makes the numbers repeat every run
rng = np.random.default_rng(0)

# a batch: 32 examples, 3 features each (random stand-in data)
X = rng.normal(size=(32, 3))
# step 1 weights: 3 inputs -> 4 outputs, stored as (outputs, inputs)
W1 = rng.normal(size=(4, 3))
b1 = np.zeros(4)                   # one bias per output, starting at 0
W2 = rng.normal(size=(2, 4))       # step 2 weights: 4 inputs -> 2 outputs
b2 = np.zeros(2)

H = X @ W1.T + b1                  # step 1
print(H.shape)                     # (32, 4)   only the last axis changed
# the bend: every negative number becomes 0
H = np.maximum(H, 0)
# element by element, so the shape is unchanged
print(H.shape)
# prints (32, 4)
Y = H @ W2.T + b2                  # step 2
print(Y.shape)                     # (32, 2)

# .size = how many numbers an array holds
n_params = W1.size + b1.size + W2.size + b2.size
print(n_params)                    # 26 = (4*3 + 4) + (2*4 + 2)

# Give the input an extra leading axis (8 groups of 128 examples): only the
# last axis changes.
X3 = rng.normal(size=(8, 128, 3))
print((X3 @ W1.T + b1).shape)      # (8, 128, 4)

All pairwise dot products in one multiplication

import numpy as np

rng = np.random.default_rng(0)
A = rng.normal(size=(4, 8))        # 4 vectors of length 8, one per row
B = rng.normal(size=(4, 8))        # another 4

S = A @ B.T                        # (4, 8) @ (8, 4) -> (4, 4)
print(S.shape)                     # (4, 4)

# Entry (i, j) is the dot product of row i of A with row j of B. Check one
# of them:
by_hand = A[1] @ B[2]              # row 1 of A dotted with row 2 of B
# (isclose allows for tiny rounding differences)
print(np.isclose(S[1, 2], by_hand))
# prints True

# One stack against itself: every pair of vectors compared in one line.
G = A @ A.T                        # (4, 4)
print(np.allclose(G, G.T))         # True: symmetric, because a·b = b·a
# square every number, then add along each row
self_dots = (A ** 2).sum(axis=1)
print(np.allclose(np.diag(G), self_dots))
# prints True: the diagonal holds each vector dotted with itself

Broadcasting, including the case that fails

import numpy as np

Z = np.zeros((4, 2))               # 4 rows, 2 columns of zeros
b = np.array([-1.0, 0.0])          # shape (2,)
# : b is stretched across the 4 rows
print((Z + b).shape)
# prints (4, 2)

col = np.array([[10.0], [20.0], [30.0], [40.0]])    # shape (4, 1)
# : the column is stretched across the 2 columns
print((Z + col).shape)
# prints (4, 2)

row_means = Z.mean(axis=1)         # one number per row
print(row_means.shape)             # (4,)
# : [:, None] adds an axis of size 1
print(row_means[:, None].shape)
# prints (4, 1)
# : now it lines up row by row
print((Z - row_means[:, None]).shape)
# prints (4, 2)

try:
    # (4, 2) and (3,): the last sizes 2 and 3 do not match
    Z + np.array([1.0, 2.0, 3.0])
except ValueError as err:
    # operands could not be broadcast together with shapes (4,2) (3,)
    print(err)

Where you will meet this again

• Week 3, penalising big weights. The exact solution of regression with a penalty is $(X^\top X + \lambda I)^{-1} X^\top \mathbf{y}$: transposes, a product, the identity matrix.
• Week 5, neural networks. A layer ? is $XW^\top + \mathbf{b}$ followed by a bend, and a network is the shape table above with more rows. PyTorch (week 6) calls the step nn.Linear and stores the weight as (outputs, inputs), exactly as here.
• Week 7, embeddings. A language model’s lookup table is $V \times d$ (50,257 × 768 for GPT-2). Looking up a token ? is reading one row, which is the same as multiplying a one-hot row of length $V$ by the matrix.
• Week 8, attention. Attention ? makes three stacks of vectors with three weight matrices, forms the $n \times n$ matrix of Fig. 5 (called $QK^\top$), turns each row into percentages, and multiplies by the third stack to get $n \times d$ again. Three matmuls, and the shape rule tells you the shape at every step. Because the middle matrix is $n \times n$, attention costs grow with the square of the text length.
• Week 10, cost tricks. The KV cache ? is two of those stacks kept in memory so they are not recomputed; its size is a shape calculation. LoRA ? adds a product of two thin matrices $BA$ to a frozen weight and uses associativity to apply it cheaply. And parameter counts are sums of $mn + m$ over the weight matrices, priced at 2 FLOPs per parameter per token.

The full interview answer

The typical form is a shape question dressed up as something else: “walk me through the shapes when a batch goes through a weight matrix” or, after week 8, “what is the shape of the attention scores and why”. Here is a complete answer in about seventy seconds.

If they push further, be ready for:

• “Why (outputs, inputs) and not (inputs, outputs)?” Convention from writing $W\mathbf{x}$ with column vectors; one row per output. With examples as rows you apply the transpose.
• “How many parameters in a layer from 768 inputs to 3072 outputs?” $768 \times 3072 + 3072 = 2{,}362{,}368$.
• “Is matmul commutative?” No. Associative and distributive, yes. Give the thin-matrices bracketing example.
• “What is the difference between * and @?” Element-wise product versus matrix product.
• “Why is attention ? quadratic in sequence length?” (Week 8.) Its score matrix is $n \times n$: $n^2 d$ work and $n^2$ memory.

Practice questions

Answer each one out loud before opening it.

Practice problems

Paper first. Write shapes before numbers, every time. Solutions are hidden under each problem.

By hand

Code

import numpy as np

def matmul(A, B):
    m, n = A.shape                  # rows and columns of A
    n2, p = B.shape                 # rows and columns of B
    if n != n2:
        raise ValueError(f"inner dimensions differ: A is {m}x{n}, B is {n2}x{p}")
    C = np.zeros((m, p))            # the answer, filled in below
    for i in range(m):              # each row of A
        for j in range(p):          # each column of B
            # the dot product of that row and that column
            for k in range(n):
                C[i, j] += A[i, k] * B[k, j]
    return C

rng = np.random.default_rng(0)
for m, n, p in [(2, 3, 2), (5, 1, 4), (3, 7, 3), (1, 6, 1)]:
    A = rng.normal(size=(m, n))
    B = rng.normal(size=(n, p))
    # same answer as NumPy, up to rounding
    assert np.allclose(matmul(A, B), A @ B)
print("all shapes agree with @")              # all shapes agree with @

try:
    matmul(np.zeros((2, 3)), np.zeros((2, 3)))
except ValueError as err:
    # inner dimensions differ: A is 2x3, B is 2x3
    print(err)

import time
import numpy as np

rng = np.random.default_rng(0)
A = rng.normal(size=(200, 200))
B = rng.normal(size=(200, 200))

start = time.perf_counter()
C_loops = matmul(A, B)                       # the function from problem 13
loop_seconds = time.perf_counter() - start

start = time.perf_counter()
C_numpy = A @ B
numpy_seconds = time.perf_counter() - start

# same answer, very different speed
assert np.allclose(C_loops, C_numpy)
numpy_ms = numpy_seconds * 1000
speedup = loop_seconds / numpy_seconds
print(f"loops: {loop_seconds:.2f} s   numpy: {numpy_ms:.2f} ms   speed-up: {speedup:,.0f}x")

n = 1000
A = rng.normal(size=(n, n))
B = rng.normal(size=(n, n))
start = time.perf_counter()
C = A @ B
seconds = time.perf_counter() - start
# n*n answers, each needing n multiplies and n adds
flops = 2 * n * n * n
# billions of operations per second
gflops_per_second = flops / seconds / 1e9
print(f"{gflops_per_second:.1f} GFLOP/s on this CPU")

# Your numbers will differ. One laptop printed:
#   loops: 5.74 s   numpy: 9.75 ms   speed-up: 588x
#   25.4 GFLOP/s on this CPU
# The ratio is the point: NumPy is hundreds to thousands of times faster
# than Python loops,
# and a GPU does tens of thousands of GFLOP/s.

import numpy as np

rng = np.random.default_rng(0)
# step 1: 3 inputs -> 4 outputs, stored as (outputs, inputs)
W1 = rng.normal(size=(4, 3)) * 0.5
b1 = np.zeros(4)
W2 = rng.normal(size=(2, 4)) * 0.5      # step 2: 4 inputs -> 2 outputs
b2 = np.zeros(2)

def forward(X):
    # every axis except the last, e.g. (32,) or (8, 128)
    leading = X.shape[:-1]
    H = X @ W1.T + b1                   # step 1
    assert H.shape == leading + (4,)
    H = np.maximum(H, 0)                # the bend: negatives become 0
    assert H.shape == leading + (4,)
    Y = H @ W2.T + b2                   # step 2
    assert Y.shape == leading + (2,)
    return Y

print(forward(rng.normal(size=(32, 3))).shape)       # (32, 2)
# : same code, extra leading axis
print(forward(rng.normal(size=(8, 128, 3))).shape)
# prints (8, 128, 2)

import numpy as np

rng = np.random.default_rng(0)
# 1,000 vectors of length 64, one per row
E = rng.normal(size=(1000, 64))
# (1000, 64) @ (64, 1000) -> (1000, 1000): every pair at once
S = E @ E.T
print(S.shape)                      # (1000, 1000)
print(S.size)                       # 1000000 dot products from one line

i, j = 3, 17
by_hand = 0.0
# the dot product of rows 3 and 17, written out
for k in range(64):
    by_hand += E[i, k] * E[j, k]
assert np.isclose(S[i, j], by_hand)
print(round(by_hand, 3))            # 5.803

# symmetric: row i with row j = row j with row i
assert np.allclose(S, S.T)
self_dots = (E ** 2).sum(axis=1)    # each vector dotted with itself
assert np.allclose(np.diag(S), self_dots)

# In the next note you divide each row by its length first;
# then this same line gives the cosine similarity of every pair.

Common mistakes

All sources

Everything linked above, in the order to use it, plus one reference.

1. 3Blue1Brown, “Linear transformations and matrices | Chapter 3” and “Matrix multiplication as composition | Chapter 4” (videos, about 10 minutes each).
2. matrixmultiplication.xyz (interactive).
3. Khan Academy, Matrices (Precalculus) (exercises). Multiplication and properties only.
4. NumPy, Broadcasting and numpy.matmul (documentation), and the site’s own NumPy reference.
5. Deisenroth, Faisal and Ong, Mathematics for Machine Learning, sections 2.2–2.3 and 2.7 (free PDF). The formal treatment of matrices and linear maps.

One-minute revision