Tensors and shapes

A tensor is an array with any number of axes. A number has none, a vector one, a matrix two; a batch of sentences in a transformer has three, (batch, tokens, features). Most of reading model code is following shapes through slicing, reshaping, broadcasting and einsum.

Key idea: Tensors generalise vectors and matrices to any number of axes; shape (B, T, d); slicing removes axes; reshape regroups without moving data; broadcasting repeats along size-1 axes; einsum names axes and sums the missing ones
Real world: A colour video is a 4-axis tensor: frames × height × width × 3 colour channels. One frame is a 3-axis slice, one pixel is a vector of 3 numbers, one brightness value is a single number.
In an LLM: GPT-2 trains on blocks of shape (batch, tokens, features), for example (8, 1024, 768): 8 texts of 1,024 tokens, 768 numbers per token. Inside attention it reshapes them to (8, 12, 1024, 64), one slice per head.

More axes

A is a block of numbers with any number of axes (also called dimensions; numpy calls the count ndim). The list of sizes along the axes is its shape:

Axes	Name	Example shape	In an LLM
0	scalar	`()`	the loss, a learning rate
1	vector	`(768,)`	one token’s hidden state
2	matrix	`(1024, 768)`	one text: a row per token
3	3-axis tensor	`(8, 1024, 768)`	a batch of 8 texts
4	4-axis tensor	`(8, 12, 1024, 1024)`	attention weights: batch, head, query, key

Nothing new happens mathematically. A tensor is a convenient way to do the same matrix operations on many matrices at once, which is exactly what GPUs are good at.

Slicing

Indexing a tensor picks out part of it. Fixing an index removes that axis; a colon : keeps it. Pick slices of a small (3, 5, 4) batch:

x[0] → shape (5, 4)

thingact.smallfirstThe0.10.01.01.0cat1.00.10.00.0sat0.21.00.00.0on0.00.30.80.0the0.10.01.00.0

axes2

numbers20

x[0] is sentence 0: a matrix, one row per token. Each index you fix removes one axis from the shape; a colon (:) keeps that axis. Fix all three and you are left with one number.

Three sentences of five tokens, four numbers per token: a tensor of shape (batch, tokens, features) = (3, 5, 4), drawn as three matrices stacked one behind the other. Pick a slice.

example · a video file

A 10-second clip at 30 frames per second, 1080 × 1920 pixels, in colour, is a tensor of shape (300, 1080, 1920, 3), about 1.9 billion numbers. video[0] is the first frame, shape (1080, 1920, 3). video[:, 540, 960] is the centre pixel’s colour over time, shape (300, 3). video[..., 0] is the red channel of everything, shape (300, 1080, 1920). Video models and image models read their inputs with exactly these slices.

Reshaping into heads

A reshape (view in PyTorch) reads the same numbers, in the same order, with a different shape. The total count must stay the same: $768 = 12 \times 64$ . A transpose of two axes swaps them. Multi-head attention uses both to give each head its own slice of every token:

(4, 6)

Thecatsaton012345012345012345012345

shape

(4, 6) 4 tokens, 6 numbers each

(4, 2, 3) split each row into 2 heads of 3

(2, 4, 3) group by head

GPT-2 small, batch B, T tokens:

(B, T, 768)

(B, T, 12, 64)

(B, 12, T, 64)

Each token is one row of 6 numbers. Press play: no number will change, only how they are grouped.

Multi-head attention splits each token’s vector into equal pieces, one per head. Here 6 numbers become 2 heads × 3. The cell labels are the original positions 0 to 5; colour shows the head.

example · 12 heads in GPT-2

GPT-2 small computes queries for all heads at once: $Q$ has shape (B, T, 768). It is viewed as (B, T, 12, 64), splitting each token’s 768 numbers into 12 groups of 64, and transposed to (B, 12, T, 64). Now each of the 12 heads has an ordinary (T, 64) matrix per text, and one batched multiplication computes every head’s scores. At the end the process runs backwards to (B, T, 768). See multi-head attention.

Broadcasting

Adding a bias of 768 numbers to a (T, 768) matrix should add it to every row. makes that automatic: line the two shapes up from the right; each pair of sizes must be equal, or one of them must be 1 (or missing), in which case that array is repeated along that axis.

(4, 3)

(3,)

(4, 3)

line the shapes up from the right

a43b3out43

A (3,) vector lines up with the last axis of (4, 3). It is repeated down all 4 rows: every token gets the same bias added. Rule: two sizes fit if they are equal or one of them is 1 (or missing).

Solid cells are stored numbers; dashed cells are copies broadcasting pretends are there. No memory is used for the copies.

Broadcasting is everywhere in model code: adding biases, scaling each token by its own norm in RMSNorm, and adding a (T, T) causal mask to scores of shape (B, 12, T, T).

Einsum: shapes as letters

(Einstein summation) writes a tensor operation by naming every axis with a letter. The rule fits in a sentence: letters that appear in the inputs but not in the output are multiplied over and summed; the others are kept. The string reads like a shape annotation, which is why model code uses it so much:

einsum('btd,de->bte')

input 1: btd (8, 10, 768)

input 2: de (768, 64)

output: bte (8, 10, 64)

keptb, t, e

summedd

multiplications3,932,160

same as: x @ W

For every sentence b and token t, sum over the 768 features d: W_Q applied to a whole batch in one go.

Each letter names an axis. Hover a letter to find it everywhere. Dotted underline: the letter is missing from the output, so that axis is summed away.

going deeper

Rules of thumb

Reshape never moves data in memory (when the tensor is contiguous); transpose changes the strides, so a following view may need .contiguous() or reshape.
Batched matrix multiply @ works on the last two axes and broadcasts over the rest: (B, 12, T, 64) @ (B, 12, 64, T) → (B, 12, T, T).
Broadcasting compares shapes from the right: (8, 1, 1024) and (1024,) fit; (8, 1024) and (8,) do not.

In code

import torch
B, T, D, H = 8, 1024, 768, 12
x = torch.randn(B, T, D)
x[0].shape, x[:, 1].shape, x[..., 0].shape     # (1024, 768), (8, 768), (8, 1024)
 
q = x.view(B, T, H, D // H).transpose(1, 2)     # (8, 12, 1024, 64)
k = q                                           # stand-in
scores = torch.einsum("bhqd,bhkd->bhqk", q, k)  # (8, 12, 1024, 1024)
same = q @ k.transpose(-2, -1)                  # identical
 
bias = torch.randn(D)
(x + bias).shape                                 # broadcast over B and T: (8, 1024, 768)
x / x.norm(dim=-1, keepdim=True)                 # (8, 1024, 1) broadcast across 768