Every tool the rest of the book leans on, from tensor shapes and the chain rule to Adam, initialization, and PyTorch internals. Built from intuition up to the depth interviewers probe.

Here is a bug that quietly ruins real training runs. You have predictions from a model and the targets you want them to match, and you compute the mean squared error. The code runs, the loss goes down, the model trains. It just trains toward the wrong thing, and nothing ever crashes.

The model's output had shape (5, 1) and the targets had shape (5,). Subtracting them did not raise an error. It quietly built a $5 \times 5$ table of every prediction minus every target, and the "loss" became the average over all 25 pairs. Most of those pairs compare example $i$ with example $j$, which is meaningless. The loss it reports is 0.30, while the true error is 0.01.

This section is about never writing that bug, and about the handful of ideas (shapes, broadcasting, memory layout, and einsum) that make tensor code readable and fast. They are the grammar of everything else in this book. Every equation from here on comes with shapes, and every interview coding round assumes you can manipulate them without thinking.

A tensor is a box of numbers with a shape

In machine learning, a tensor is an $n$-dimensional array of numbers that all share one type. The word comes from physics, where it means something more specific, but in code it means exactly what NumPy calls an ndarray and PyTorch calls a Tensor. Three properties describe it.

• The shape is a tuple of sizes, one per axis. An array of shape (2, 3, 4) has 24 numbers.
• The rank (or ndim) is the number of axes. A scalar has rank 0, a vector rank 1, a matrix rank 2.
• The dtype is the number format, such as float32, bfloat16, or int64. It sets the memory per element and the precision, which matters a great deal in §0.9.

Almost every tensor you meet in modern AI has a conventional shape, and knowing the conventions lets you read code at a glance. These are the ones this book uses.

Here $B$ is the batch size, $T$ the sequence length, $d$ the model width, $h$ the number of attention heads, and $C, H, W$ the channels, height, and width of an image. That last pair of rows trips people up constantly. PyTorch stores a linear layer's weight as (out, in), so when you port a NumPy derivation to PyTorch you will see a transpose appear. Neither is wrong. Just know which one you are looking at.

How big is a tensor

A tensor's memory is its number of elements times the bytes per element, 4 for float32 and 2 for bfloat16. That one multiplication tells you, before you run anything, which tensors in a model are cheap and which will not fit. Take GPT-2 small ($d = 768$, $h = 12$ heads, 12 layers, a vocabulary of 50,257 tokens) and a modest batch of $B = 8$ sequences of $T = 1{,}024$ tokens.

The hidden states are small. The attention scores, one $T \times T$ table per head, are about 30 times larger per layer, and storing them for every layer costs more than the model itself. The logits are a surprise to most people, because a 50,257-way score for every token outweighs the whole model. Those two numbers motivate two pieces of engineering you will meet later, FlashAttention, which never stores the score table (§12.3), and computing the loss over the vocabulary in chunks. Get in the habit of doing this arithmetic, because "will it fit" is the first question of every systems interview.

Reductions remove the axis you reduce over

A reduction (sum, mean, max, and so on) collapses one axis. If X has shape (B, d), then X.sum(axis=0) adds down the batch and leaves shape (d,), while X.sum(axis=1) adds across features and leaves (B,). The axis you name is the axis that disappears.

Often you want to keep a size-1 placeholder where the axis was, so that the result lines up with the original for the next operation. That is what keepdims=True does.

That last line normalizes each row by its own total. Without keepdims, the divisor would have shape (2,), and dividing a (2, 3) array by it raises an error here only because 2 does not equal 3. With a square matrix it would not raise anything. It would divide each column by a row total, which is the same silent bug as at the top of the page. Karpathy hits exactly this while normalizing a count matrix in the makemore video linked below.

Broadcasting, the rule that saves loops and causes bugs

Broadcasting is how NumPy and PyTorch combine arrays of different shapes in elementwise operations. You have already used it if you have ever written X - X.mean(axis=0) or logits + bias. The rule is short, and it is worth knowing exactly.

To combine shapes $a$ and $b$, line them up from the right. If one shape has fewer axes, pad it on the left with 1s until both have the same length. Then compare the shapes axis by axis.

1. If the two sizes are equal, that is the output size.
2. If one of them is 1, the output size is the other one, and the size-1 array is repeated (stretched) along that axis.
3. Otherwise the shapes are incompatible and you get an error.

In plain words, trailing axes must match, and a size of 1 (or a missing axis on the left) means "use the same value for every position".

Let's apply it to the bug. The prediction had shape (5, 1) and the target (5,). Line them up from the right. The target becomes (1, 5) after padding. Axis by axis, $5$ against $1$ gives $5$, and $1$ against $5$ gives $5$. The result is (5, 5), which is what we saw.

Try these and your own in the visualizer. Click a dimension to change it, and add or remove axes.

Making broadcasting do what you mean

Most of the time you choose which axes line up by inserting size-1 axes yourself. Indexing with None (also spelled np.newaxis) inserts one, so for u of shape (3,), u[:, None] is a (3, 1) column and u[None, :] a (1, 3) row. Then u[:, None] - v[None, :] is the table of every difference $u_i - v_j$, on purpose this time. The reverse operation, squeeze, removes size-1 axes, and it is the fix for the bug at the top of the page, y_pred[:, 0] or y_pred.squeeze(-1). PyTorch calls the same two operations unsqueeze and squeeze.

Broadcasting never copies the data

When NumPy "repeats" the size-1 array, it does not actually build the bigger array. It reads the same memory again and again. To see how, we need to look at how arrays are laid out in memory, which also explains views, copies, and why transposes are free.

Strides, views, and copies

Memory is a flat line of bytes. An array of shape (2, 3, 4) stores its 24 numbers in one contiguous block, and a small tuple called the strides says how many bytes to jump to move one step along each axis. For float64 (8 bytes per number) in the default row-major (C) order, the last axis is contiguous, so its stride is 8. One step along the middle axis skips a whole row of 4 numbers, 32 bytes. One step along the first axis skips $3 \times 4$ numbers, 96 bytes.

The position of element $(i_1, \dots, i_n)$, in bytes from the start of the buffer, is

where $s_k$ is the stride of axis $k$ in bytes. In plain words, to find an element, walk $i_k$ steps of size $s_k$ along each axis and add up the jumps.

This one formula explains a lot of tensor behavior.

• Transpose is free. Swapping two axes swaps their strides and touches no data. The result is a view, a new shape and stride tuple pointing at the same memory.
• Basic slicing is free. X[:, ::2] doubles a stride. Still a view.
• Broadcasting is free. A stretched axis gets stride 0, so every step along it reads the same numbers.
• Some operations must copy. Fancy indexing with an integer array (X[[0, 2]]) gathers arbitrary rows, which no stride tuple can describe, so it copies. Reshaping a transposed array usually copies too, as we'll see.

A contiguous array is one whose elements sit in memory in exactly the order you would visit them in a row-major loop. A transposed array is not contiguous. That matters for speed, because GPUs and CPUs read memory fastest in long consecutive runs (you'll meet this as memory coalescing in §11.3). It also matters for correctness in PyTorch, where x.view(...) refuses any new shape that the existing strides cannot express (as with most reshapes of a transposed tensor), and you must call x.contiguous() first or use x.reshape(...), which copies when it has to.

def broadcast_shape(s1, s2):
    """s1, s2: tuples of ints -> broadcast shape (tuple), or raise ValueError"""
    raise NotImplementedError

Pairwise distances, the classic broadcasting problem

A problem you will meet in k-means, in nearest-neighbor search, and in interviews is to compute every squared distance between the rows of $\mathbf{A} \in \mathbb{R}^{n \times d}$ and the rows of $\mathbf{B} \in \mathbb{R}^{m \times d}$, giving $\mathbf{D} \in \mathbb{R}^{n \times m}$ with $D_{ij} = \lVert \mathbf{a}_i - \mathbf{b}_j \rVert^2$.

There are two good ways. The direct one is a broadcast, A[:, None, :] - B[None, :, :], of shape $(n, m, d)$, squared and summed over the last axis. It is clear, but it builds an $n \times m \times d$ intermediate. With $n = m = 10{,}000$ and $d = 768$ that is $7.7 \times 10^{10}$ numbers, about 300 GB in float32.

The better one expands the square (the same identity that ties distance to the dot product for word vectors in §1.3),

In plain words, the squared distance is two squared norms minus twice a dot product, and all $nm$ dot products together are one matrix multiply, $\mathbf{A}\mathbf{B}^\top$. The norms are an $(n, 1)$ column and a $(1, m)$ row, and broadcasting adds them into an $(n, m)$ table. Memory drops from $O(nmd)$ to $O(nm)$, and the work goes to a matmul, which is the fastest operation on any accelerator.

import numpy as np

def pairwise_sq_dists(A, B):
    """A: (n, d), B: (m, d) -> (n, m) squared distances"""
    raise NotImplementedError

Reshape is not transpose

Reshape and transpose both change a shape, and that is where the resemblance ends. Reshape keeps the elements in the same row-major order and only regroups them into new axes. Transpose (or permute in PyTorch) reorders the axes, which changes which elements are neighbors. Confusing the two produces tensors of exactly the right shape and completely wrong content, the most dangerous kind of bug.

The canonical example is splitting a model's hidden state into attention heads, which you will do in multi-head attention (§4.5). Take one sequence of $T = 2$ tokens with width $d = 4$, and split it into $h = 2$ heads of width $d/h = 2$. Each head should see both tokens, restricted to its own slice of features. Head 0 gets features 0 and 1 of every token, and head 1 gets features 2 and 3.

The right version gives head 0 the rows [0 1] and [4 5], the first two features of each token. The wrong version gives head 0 all of token 0 and head 1 all of token 1. Both have shape (1, 2, 2, 2), both train, and the wrong one is a different (worse) model. The rule of thumb is that you may only reshape to split or merge axes that are adjacent and in the right order. Here we want to split $d$ into $(h, d/h)$, which is adjacent, and then move $h$ in front of $T$, which is a reordering and needs a transpose.

import numpy as np

def split_heads(X, h):
    raise NotImplementedError

def merge_heads(Y):
    raise NotImplementedError

Sliding windows with strides alone

Strides can do more than describe transposes. Any regular pattern of reads can be described by a shape and a stride tuple, including overlapping ones. A sliding window of width $k$, moving step elements at a time over a 1-D signal of length $n$, is an array of shape $(n_{\text{win}}, k)$ with $n_{\text{win}} = \lfloor (n - k)/\text{step} \rfloor + 1$ windows. Its row stride is step elements and its column stride is one element. No data is copied, even though windows overlap. This is the idea behind fast convolution in §0.8, and numpy.lib.stride_tricks.sliding_window_view uses exactly this.

import numpy as np
from numpy.lib.stride_tricks import as_strided

def sliding_windows(x, k, step=1):
    raise NotImplementedError

Einsum, the universal contraction

Matrix multiply, batched matmul, dot products, outer products, traces, transposes, and sums are all instances of one operation. You multiply some tensors elementwise along matching indices and then sum over the indices you don't want to keep. Einstein summation notation, einsum, writes that operation directly.

An einsum string names each axis of each input with a letter, then names the axes of the output after an arrow. The rule is that any letter that appears in the inputs but not in the output is summed over. Letters shared between inputs are multiplied together along that axis.

In plain words, $j$ appears in both inputs and not in the output, so it is the summed (contracted) index. That is a matrix multiply.

The last two rows are the whole of attention's matrix arithmetic, which is why einsum shows up all over research code. The cost of an einsum is easy to read off too. It is the product of the sizes of all the distinct letters, because that is how many multiply-adds the nested loops would do. For bhqd,bhkd->bhqk that is $B \cdot h \cdot T \cdot T \cdot d_h$, the famous quadratic cost of attention.

Contraction order changes the cost, not the answer

When you multiply three or more tensors, the result does not depend on the order you contract them in, but the cost can differ by orders of magnitude. Take $\mathbf{A} \in \mathbb{R}^{1000 \times 10}$, $\mathbf{B} \in \mathbb{R}^{10 \times 1000}$, and $\mathbf{C} \in \mathbb{R}^{1000 \times 10}$. Computing $(\mathbf{A}\mathbf{B})\mathbf{C}$ first builds a $1000 \times 1000$ matrix ($10^7$ multiply-adds), then multiplies it by $\mathbf{C}$ (another $10^7$), for $2 \times 10^7$ in total. Computing $\mathbf{A}(\mathbf{B}\mathbf{C})$ first builds a $10 \times 10$ matrix ($10^5$) and then a $1000 \times 10$ result ($10^5$), for $2 \times 10^5$. That is 100 times less work for the same answer. np.einsum(..., optimize=True) searches for a good order automatically. The same idea, applied by hand, is how LoRA keeps its cost low in §14.5, by never forming the product of its two thin matrices.

import numpy as np

def batched_matmul(A, B):
    """A: (b, n, k), B: (b, k, m) -> (b, n, m)"""
    raise NotImplementedError

def attention_scores(Q, K):
    """Q: (b, h, tq, d), K: (b, h, tk, d) -> (b, h, tq, tk), unscaled"""

Batched matrix multiply

The @ operator (np.matmul, torch.matmul) treats the last two axes of each operand as matrices and broadcasts over all the others. So (B, T, d) @ (d, e) gives (B, T, e), a linear layer applied to every token at once, because the weight is treated as a single matrix broadcast over $B$. And (B, h, T, d_h) @ (B, h, d_h, T) gives (B, h, T, T), one score table per head per sequence. Two special cases are worth remembering. A 1-D operand on the left is treated as a row vector and on the right as a column vector, and the added axis is removed from the result. That is why v @ M and M @ v both work without reshaping.

Underneath, batched matmul is one call to a highly tuned kernel (a batched GEMM, for "general matrix multiply"). A Python loop over the batch would launch $B$ separate small kernels, which on a GPU is far slower. Writing operations in batched form is the single most important performance habit in this book. You'll see why at the hardware level in the roofline model (§11.2). This section treats a matrix as a box of numbers to be moved and multiplied. What a matrix does to space, its eigenvalues, singular values, and norms, is the subject of §0.10.

The shape bugs that never crash

Here are the bugs that most often survive into real training runs, in rough order of how often they bite.

The common thread is that broadcasting and reshaping are permissive by design, so the type system will not save you. Only shape comments, explicit axes, and assertions will.

In PyTorch

PyTorch tensors follow the same rules, with a few extra names. view is a reshape that never copies (and fails if it would need to), reshape copies when it must, permute is a general transpose, expand is a zero-copy broadcast (stride 0), and repeat really copies.

Two PyTorch details are worth a sentence each. Strides in PyTorch are counted in elements, not bytes. And x.T on tensors that are not 2-D is deprecated, because reversing every axis is rarely what you want. Use x.mT to swap the last two axes, or permute for anything else.

Read more

• NumPy documentation. Broadcasting. The official statement of the rule, with pictures.
• NumPy documentation. Internal memory layout of an ndarray. Strides and contiguity, precisely.
• Harris et al. (2020). Array programming with NumPy. The design of the array model that PyTorch and JAX inherited, with clear figures of strides and broadcasting.
• PyTorch documentation. Tensor Views. Which operations return views and which copy.
• Edward Z. Yang (2019). PyTorch internals. How a tensor is a storage plus sizes, strides, and an offset, from one of PyTorch's core developers.
• Tim Rocktäschel (2018). Einsum is All you Need - Einstein Summation in Deep Learning. A short, readable tour of einsum for deep learning.

Watch next

▶

▶