Tensors
Arrays with more axes: batches, sequences, heads.
Create a free account Try part of it here
Why it matters
Real AI models never look at one example at a time. A language model works on a batch of sentences at once. Each sentence is a row of words, and each word is a list of numbers. That's three layers of grouping. Arrays with layers like this are called tensors (the main data type in PyTorch is literally called Tensor), and most of the code inside a model is reshaping them and multiplying them. If you can follow the shapes, you can follow the model.
What you'll be able to do
- Describe scalars, vectors, matrices and tensors by their number of axes
- Read shapes like (batch, tokens, features) and say what each axis means
- Reduce, reshape and reorder axes with NumPy
- Use batched matrix multiplication the way transformers do
What it covers
- Boxes inside boxes
- Reading a tensor's shape
- Averaging along an axis
- Reshaping: same numbers, new grouping
- Multiplying stacks of tables
- Name your axes
- Check shapes often
Boxes inside boxes
Think about how a filing cabinet is organized:
- one number is a single fact;
- a list of numbers is one form (a vector);
- a table of numbers is a folder of forms (a matrix);
- a stack of tables is a drawer of folders;
- a stack of stacks is a whole cabinet.
A tensor is any of these: numbers grouped in as many layers as you need. Each layer of grouping is called an axis. A list has one axis, a table has two, a stack of tables has three.
In plain words: a tensor is an array with any number of axes. Its shape lists how big each axis is, from the outside in.
Reading a tensor's shape
The shape is the key to everything. Here's how the shapes inside a language model read:
| Axes | Everyday name | Example shape | What it holds |
|---|---|---|---|
| 0 | a single number | () | one score |
| 1 | a list (vector) | (768,) | one word's 768 numbers |
| 2 | a table (matrix) | (128, 768) | a sentence: 128 words × 768 numbers |
| 3 | a stack of tables | (32, 128, 768) | 32 sentences at once |
| 4 | a stack of stacks | (32, 12, 128, 128) | attention scores: sentences × heads × words × words |
Photos work the same way: a batch of color photos is (photos, height, width, 3 color values).
One habit makes shapes easy to read: the first axis is usually the batch (how many examples) and the last axis is usually the features (the numbers describing each thing).
import numpy as np batch = np.zeros((32, 128, 768)) print(batch.ndim, batch.shape, batch.size) # axes, shape, total numbers print(batch[0].shape) # one sentence print(batch[0, 5].shape) # one token of one sentence print(batch[:, 0].shape) # the first token of every sentence
Picking with a number removes that axis: batch[0] is one sentence, a table. Picking with : keeps the axis: batch[:, 0] is the first word of every sentence.
batch = [[[0] * 4 for _ in range(3)] for _ in range(2)] print(len(batch), len(batch[0]), len(batch[0][0]))
Show the answer
2 3 4
Read shapes from the outside in: 2 tables, each with 3 rows of 4 numbers.
Show the answer
16
The first axis is the batch: 16 photos, each 224 × 224 pixels with 3 color values.
Averaging along an axis
The axis setting from the matrices lesson works on any tensor. Say you want one vector for a whole sentence instead of one per word. Average over the word axis:
import numpy as np rng = np.random.default_rng(0) batch = rng.normal(size=(4, 6, 8)) # 4 sentences, 6 tokens each, 8 features sentence_vectors = batch.mean(axis=1) # average over tokens print(sentence_vectors.shape) # (4, 8): one vector per sentence print(batch.max(axis=-1).shape) # axis=-1 means "the last axis"
The word axis (axis 1) disappeared: 4 sentences, 8 numbers each. This is called mean pooling, and it's a common way to get one embedding for a whole document. axis=-1 is a handy shortcut for "the last axis".
In plain words: averaging along an axis squashes that axis away. Shape (4, 6, 8) averaged over axis 1 becomes (4, 8).
A table is a tensor of shape (2, 3). Sum along each axis and see which dimension disappears.
Show the answer
(16, 512)
Averaging over axis 1 (tokens) removes it: one 512-number vector per document.
Reshaping: same numbers, new grouping
reshape regroups the same numbers into a different shape, like re-filing the same forms into different folders. The total count of numbers must stay the same. transpose with a list of axes changes their order.
Attention uses both to split each word's numbers into several heads: smaller groups that each look for a different kind of relationship.
import numpy as np B, T, H, Dh = 2, 5, 4, 3 # batch, tokens, heads, features per head x = np.arange(B * T * H * Dh).reshape(B, T, H * Dh) print(x.shape) # (2, 5, 12) heads = x.reshape(B, T, H, Dh).transpose(0, 2, 1, 3) print(heads.shape) # (2, 4, 5, 3): batch, heads, tokens, features
Each word's 12 numbers are split into 4 heads of 3. transpose(0, 2, 1, 3) means "new order: old axis 0, then 2, then 1, then 3", which moves heads in front of words.
Multiplying stacks of tables
When you use @ on tensors with more than two axes, NumPy multiplies the last two axes as tables and repeats that for everything in front. So one line computes the attention table for every head of every sentence:
import numpy as np rng = np.random.default_rng(1) B, H, T, Dh = 2, 4, 5, 3 Q = rng.normal(size=(B, H, T, Dh)) K = rng.normal(size=(B, H, T, Dh)) scores = Q @ K.transpose(0, 1, 3, 2) # swap the last two axes of K print(scores.shape) # every token vs every token, per head, per sentence print(np.allclose(scores[1, 2], Q[1, 2] @ K[1, 2].T))
Only the last two axes take part in the multiplication: (5 × 3) @ (3 × 5) = (5 × 5). The first two (2 sentences, 4 heads) just come along for the ride. The last line checks one of those 8 tables by hand.
Go deeper The shape bookkeeping
Read it as: "batch, heads, tokens, head-features, times the same with the last two axes swapped, gives batch, heads, tokens, tokens."
Sometimes you need an axis of size 1 so shapes line up for broadcasting. keepdims=True keeps the averaged axis as size 1 instead of removing it:
import numpy as np x = np.arange(12, dtype=float).reshape(2, 3, 2) means = x.mean(axis=1, keepdims=True) # shape (2, 1, 2), not (2, 2) print(means.shape, (x - means).mean(axis=1).round(6).tolist())
Show the answer
Exactly 24 numbers in total
Reshape regroups the same numbers, so the total must match: 6 × 4 = 24.
Name your axes
Tensors get confusing because the numbers in a shape don't say what they mean. (8, 256, 512) could be anything. So name the axes as you go: 8 texts, 256 tokens each, 512 features per token. Good code does this in comments. You can hardly debug a shape error without it.
There's nothing exotic here, by the way. "Tensor" is the word for an array with any number of axes. A list of numbers is a tensor, and so is a table.
reshape and transpose aren't the same
Both change the shape, so they look alike. They do very different things.
import numpy as np x = np.arange(6).reshape(2, 3) print(x) print(x.reshape(3, 2)) # same reading order, new grouping print(x.T) # rows and columns swapped
Both results have shape (3, 2), but the numbers land in different places. reshape keeps the reading order and regroups it. A transpose reorders the axes. Using the wrong one raises no error. You get the right shape holding scrambled data, and a model that quietly fails to learn. That's among the hardest bugs in machine learning to find, so stop and ask which one you mean.
One last trap: taking an average removes an axis, and that can break the next step. keepdims=True keeps it in place with a size of 1, so that broadcasting still lines up.
Check shapes often
Printing .shape after each step takes a second and catches most mistakes early.
Practice what you just read
- 6 code examples you can edit and run
- a practice exercise that checks your code
- a 5-question quiz
- 4 review cards for spaced repetition
Try a quiz question
What does @ do with two 4D tensors of shapes (B, H, T, D) and (B, H, D, T)?