1.24 Tensors and einsum: study notes

Learn AI by Story

1. The idea in one sentence

A tensor is just a box of numbers with any number of directions (axes), and einsum is a tiny labelling language that lets you say "multiply these boxes, line up the matching labels, add up the repeated ones" in a single line.

2. The story from the video

Mei takes a selfie. Her phone does not store a face; it stores a grid of tiny coloured dots, and each dot is three numbers (how much red, green and blue). Stack the rows, the columns and the three colours and you get one box of numbers: 224 tall, 224 wide, 3 deep. That is 224 x 224 x 3 = 150,528 numbers for one small photo. When the app looks at 32 selfies at once, the box gets one more direction, and now it holds 4,816,896 numbers. Every one of those boxes is a tensor, and einsum is how we tell the computer how to combine them without writing loops.

3. Key definitions (glossary)

Scalar: a single number, like 5 or 0.37. It has no axes. Shape: ().
Vector: a list of numbers, one axis. Example (1, 2, 3) has shape (3,).
Matrix: a grid of numbers, two axes (rows and columns). A 2 x 3 matrix has shape (2, 3).
Tensor: the general word for an array with any number of axes. Scalars, vectors and matrices are tensors with 0, 1 and 2 axes. Goodfellow's Deep Learning describes tensors as arrays with a variable number of axes.
Rank / order of a tensor: the number of axes. A photo batch has rank 4. Warning, two meanings of "rank": for a matrix, "rank" in linear algebra means the number of independent rows or columns (a measure of how much information the matrix carries, see 1.21 basis). A 3 x 3 matrix of all ones has tensor rank 2 (two axes) but matrix rank 1. Deep learning libraries call the number of axes ndim; many textbooks prefer the word "order" for that to avoid the clash.
Shape: the list of sizes along each axis, such as (32, 224, 224, 3). The product of the shape is the total count of numbers.
Axis / dimension: one direction of the box. Careful: "dimension" is also used for the length of a vector (a 3-dimensional vector has one axis of size 3). In code, "axis 0" means the first direction.
Reshape: re-reading the same numbers with a different shape, e.g. 6 numbers as (2, 3) or (3, 2) or (6,). The count never changes.
Broadcasting: an automatic rule that stretches a small tensor (for example a bias vector) across a bigger one so you can add or multiply them without copying.
Reduction: collapsing one or more axes by summing, averaging, taking a max, and so on. Summing each row of a matrix reduces the column axis.
Contraction: a reduction that happens while two tensors are multiplied: multiply matching entries along a shared axis, then add them. The dot product and matrix multiplication are contractions.
Outer product: the opposite of a contraction. Multiply every entry of one vector by every entry of another to build a grid, with nothing added up.
Einstein summation (einsum): a notation (from physics) where each axis gets a letter, and any letter that is repeated between inputs and missing from the output is summed over.
Implicit vs explicit mode: in implicit mode you write only the inputs, like 'ij,jk', and the library decides the output (letters appearing once, sorted alphabetically). In explicit mode you add -> and name the output yourself, like 'ij,jk->ik'. Explicit mode is clearer and lets you do things implicit mode cannot, such as keeping a repeated letter or reordering output axes.
NCHW vs NHWC: two conventions for the order of image axes. N = number of images (batch), C = channels (colours), H = height, W = width. PyTorch usually uses NCHW; TensorFlow/Keras usually default to NHWC. The numbers are the same, only the axis order differs, and mixing them up gives silent wrong answers.

4. Detailed step-by-step explanation

Step 1: from a selfie to a box of numbers

A greyscale picture is a matrix: each cell is a brightness from 0 (black) to 255 (white). A colour picture needs three such matrices, one each for red, green and blue, so we get height x width x 3. For a 224 x 224 image, that is 224 x 224 x 3 = 150,528 numbers. Nothing mystical: it is a long list that we agreed to arrange in a box so the neighbours stay neighbours.

Step 2: why add more axes?

GPUs are fast when they do the same thing to lots of data at once. So we stack 32 images into one box: shape (32, 224, 224, 3), which is 32 x 150,528 = 4,816,896 numbers. The first axis is the "batch" axis. Operations that treat each image independently simply act along the other axes.

Step 3: the shape tells the story

When you see a tensor in code, first read its shape. Ask: which axis is the batch, which is position, which is feature? Almost every bug in deep learning code is a shape bug, and most can be caught by writing the shape of each variable in a comment.

Step 4: the pain of loops

To multiply two matrices by hand you write three nested loops: for each row i, for each column k, add up over j. For batched data you add a fourth loop. It works, but it is slow in plain Python and easy to get wrong.

Step 5: labelling the axes

Einsum says: give every axis a letter. For matrix multiplication of A (shape i x j) with B (shape j x k), write 'ij,jk->ik'. Read it aloud: "A has axes i and j, B has axes j and k, the answer has axes i and k." The letter j appears in both inputs but not in the output, so it is multiplied along and summed away. That is exactly matrix multiplication.

Step 6: the three rules

Same letter on different inputs means: those axes must have the same length, and entries are matched up.
A letter that is missing from the output is summed over.
A letter that is in the output stays as an axis, in the order you write it.

That is the whole language. Everything else is a consequence.

Step 7: a tour of one-liners

'i,i->' dot product: (1,2,3) and (4,5,6) give 1x4 + 2x5 + 3x6 = 32.
'ii->' trace: sum of the diagonal. For [[1,2],[3,4]] it is 1 + 4 = 5.
'ij->ji' transpose: swap which letter is the row.
'ij,jk->ik' matrix multiplication.
'i,j->ij' outer product.
'ij->i' sum across each row.
'bij,bjk->bik' batch matrix multiplication: b is the batch letter, so each batch item gets its own matrix product.
'bqd,bkd->bqk' attention scores: for each batch item b, take query vectors (q positions, d features) and key vectors (k positions, d features), and take the dot product over d. The result says how much each query looks at each key.

Step 8: why einsum matters

It is readable (the letters document the shapes), it is short, and the library can pick an efficient computation order (NumPy has einsum_path for this). It also removes the need to remember when to call transpose before a matmul.

5. Everything the video skipped

Tensor rank, again. Because of the clash in meanings, if someone says "rank 2" check the context. A matrix of shape (1000, 1000) has two axes (tensor order 2) but could have matrix rank as low as 1. Low-rank matrices matter in AI: methods such as LoRA fine-tuning store a big update as the product of two thin matrices.

Contiguous memory and reshape. A tensor lives in memory as one long line of numbers. Reshape only changes how we index that line, so it is usually free. Transpose also usually just changes the indexing rule (strides) instead of moving data, which is why some libraries need an extra contiguous step before some operations.

Broadcasting rules. Libraries compare shapes from the last axis backwards. Two sizes are compatible if they are equal or one of them is 1 (or missing). Adding a bias of shape (3,) to an image batch of shape (32, 224, 224, 3) works because the last axes match. Adding a bias of shape (3,) to shape (32, 3, 224, 224) does not do what you hope. Always check that the right axes line up.

Reductions. Sum, mean, max along an axis remove that axis (unless you ask to keep it). The mean over the batch axis gives the "average image". Reductions are what loss functions do: they compress many per-example numbers into one.

Contraction vs reduction vs product. The big family tree: outer product creates axes; elementwise multiply keeps axes; contraction removes shared axes. Einsum covers all three.

Implicit mode traps. In implicit mode 'ij,jk' the output is ordered alphabetically, and 'ji' alone would transpose by sorting. This is convenient but it surprises beginners. In explicit mode you control everything, and the NumPy documentation shows that some operations (like keeping a repeated label) need explicit mode.

Ellipsis. You can write '...ij,...jk->...ik' to mean "any number of leading batch axes". This makes code work for rank 2, 3 or 4 tensors alike.

Related functions. matmul or the @ operator, tensordot, dot, outer, trace. Einsum can reproduce all of them. In practice matmul is sometimes faster for plain matrix products because it calls highly tuned libraries, so einsum shines for unusual contractions, not necessarily for plain ones.

Limits. Einsum letters must be single characters (52 letters in NumPy, upper and lower), the result can be slow if you contract many tensors in a bad order, and a wrong letter gives a shape error or, worse, a valid but wrong answer. It does not check meaning, only sizes.

Memory is the real cost. 4,816,896 numbers stored as 32-bit floats is about 19 MB for one batch of 32 small images. Bigger models and bigger images multiply that quickly, which is why batch size is limited by GPU memory.

Where this sits in the roadmap. 1.18 linear transformations showed that a matrix is a machine that moves vectors; 1.19 determinant and 1.21 basis explained how to think of what a matrix does and carries. Tensors extend this to many axes. Next, 1.25 Jacobian and Hessian uses matrices of derivatives (and the Hessian of a vector function is a rank-3 tensor), and 1.26 Taylor approximation builds on them. Later, backpropagation will chain these tensors together, and attention in transformers is einsum 'bqd,bkd->bqk' followed by a softmax.

6. Worked examples

Example A: the video's matrix multiplication

Let A = [[1, 2], [3, 4]] and B = [[5, 6], [7, 8]]. Compute 'ij,jk->ik'.

Entry (0,0): sum over j of A[0,j] x B[j,0] = 1x5 + 2x7 = 5 + 14 = 19.
Entry (0,1): 1x6 + 2x8 = 6 + 16 = 22.
Entry (1,0): 3x5 + 4x7 = 15 + 28 = 43.
Entry (1,1): 3x6 + 4x8 = 18 + 32 = 50.

Result: [[19, 22], [43, 50]]. Trace of A with 'ii->' is 1 + 4 = 5. Transpose with 'ij->ji' gives [[1, 3], [2, 4]]. Dot product of (1,2,3) and (4,5,6) with 'i,i->' is 4 + 10 + 18 = 32.

Example B: counting numbers in the selfie batch

224 x 224 = 50,176 pixels. Times 3 colours = 150,528 numbers. Times 32 images = 4,816,896 numbers. In NCHW the shape is (32, 3, 224, 224); in NHWC it is (32, 224, 224, 3). Same count, different order.

Example C (new): outer product, reshape and row sums

Take a = (1, 2, 3) and b = (10, 20). With 'i,j->ij' every entry of a meets every entry of b:

	10	20
1	10	20
2	20	40
3	30	60

No letter is repeated, so nothing is summed, and the shape is (3, 2).

Now M = [[1, 2, 3], [4, 5, 6]] with shape (2, 3), six numbers in total.

'ij->i' sums over j: row sums (1+2+3, 4+5+6) = (6, 15).
'ij->j' sums over i: column sums (1+4, 2+5, 3+6) = (5, 7, 9).
'ij->' sums everything: 21 (check: 6 + 15 = 21 and 5 + 7 + 9 = 21).
Reshape M to (3, 2): [[1, 2], [3, 4], [5, 6]]. Reshape to (6,): (1, 2, 3, 4, 5, 6). The numbers are read in the same order each time.

Example D (new): a tiny batch matrix multiply

Two items in the batch. Q = [ [[1,0],[0,1]], [[0,1],[1,0]] ] (an identity matrix and a swap matrix) and X = [ [[1,2],[3,4]], [[5,6],[7,8]] ]. Compute 'bij,bjk->bik'.

Batch 0: identity times [[1,2],[3,4]] leaves it unchanged: [[1,2],[3,4]].
Batch 1: the swap matrix swaps the two rows of [[5,6],[7,8]]: [[7,8],[5,6]].

The output has shape (2, 2, 2): [ [[1,2],[3,4]], [[7,8],[5,6]] ]. These values were checked by running the code in NumPy.

Code snippet

import numpy as np
A = np.array([[1, 2], [3, 4]])
B = np.array([[5, 6], [7, 8]])
print(np.einsum('ij,jk->ik', A, B))   # [[19 22] [43 50]]
print(np.einsum('ii->', A))            # 5
print(np.einsum('ij->ji', A))          # transpose
print(np.einsum('i,i->', [1,2,3], [4,5,6]))  # 32
print(np.einsum('i,j->ij', [1,2,3], [10,20]))  # outer product

7. Comparison tables

Objects by number of axes

Object	Axes (order)	Example shape	Example
Scalar	0	()	a loss value 0.37
Vector	1	(3,)	an embedding, (1, 2, 3)
Matrix	2	(2, 3)	weights of one layer
3-tensor	3	(224, 224, 3)	one colour image
4-tensor	4	(32, 224, 224, 3)	a batch of images

Einsum cheat sheet

Task	Explicit string	Letters summed?
Dot product	`i,i->`	i
Trace	`ii->`	i
Transpose	`ij->ji`	none
Row sums	`ij->i`	j
Outer product	`i,j->ij`	none
Matrix product	`ij,jk->ik`	j
Batch matrix product	`bij,bjk->bik`	j
Attention scores	`bqd,bkd->bqk`	d

Two image layouts

Layout	Order	Typical home	Shape for 32 selfies
NCHW	batch, channels, height, width	PyTorch (default)	(32, 3, 224, 224)
NHWC	batch, height, width, channels	TensorFlow / Keras (default)	(32, 224, 224, 3)

Two meanings of rank

Meaning	Counts what	Matrix of all ones, 3 x 3
Tensor rank / order	number of axes	2
Matrix rank	independent rows/columns	1

8. Common mistakes and myths

Myth: a tensor is something scary from physics. In deep learning, "tensor" almost always just means "array with some axes". The deeper physics meaning (how it transforms under change of coordinates) is rarely needed here.
Mistake: mixing up the two meanings of rank. Ask yourself whether the speaker is counting axes or independent columns.
Mistake: mixing NCHW and NHWC. A model trained with one layout and fed the other may run without error and give nonsense.
Mistake: reshape instead of transpose. Reshaping (6,) into (2, 3) is not the same as transposing. Reshape keeps the reading order; transpose swaps axes. Using reshape to "swap" axes scrambles pixels.
Mistake: forgetting that a repeated letter is summed. If you want to keep an axis, put the letter in the output after the arrow.
Mistake: relying on implicit mode. Alphabetical output order silently changes results. Prefer ->.
Myth: einsum is always faster. It is clearer; speed depends on the contraction order and the backend.
Mistake: assuming broadcasting does what you mean. It matches from the last axis, so a mis-shaped bias still broadcasts and gives wrong, not erroring, answers.
Myth: more axes means a smarter model. Axes are bookkeeping for the data. Intelligence comes from the learned numbers.

9. Real-life and AI uses

Images and video: photos (height, width, channels), batches of photos, video with an extra time axis.
Text: a batch of sentences becomes (batch, sequence length, embedding size).
Attention in transformers: 'bqd,bkd->bqk' computes how strongly each word looks at every other word, for every sentence in the batch, in one line.
Multi-head attention: adds a head axis, (batch, head, position, feature), and einsum handles it without extra loops.
Convolutions: filters are 4-axis tensors (output channels, input channels, height, width).
Physics and graphics: stress tensors, 3D transformations.
Probabilistic models: sums over many variables (tensor networks, graphical models) are written as einsums.
Audio: spectrograms are (batch, frequency, time).

10. Self-check questions

How many axes does a scalar have? Zero.
What is the shape of one 224 x 224 colour image in NHWC? (224, 224, 3).
How many numbers are in a batch of 32 such images? 32 x 224 x 224 x 3 = 4,816,896.
What does 'ij,jk->ik' compute? Matrix multiplication; j is summed.
What is np.einsum('ii->', [[1,2],[3,4]])? The trace, 1 + 4 = 5.
What are the two meanings of "rank"? Number of axes of a tensor, versus the number of independent rows/columns of a matrix.
What is the shape of the result of 'i,j->ij' for lengths 3 and 2? (3, 2).
In 'bqd,bkd->bqk', which letter is summed and which is the batch? d is summed; b is the batch.
Why might reshape be wrong when you really need transpose? Reshape keeps memory order, transpose swaps axes.
What is the difference between implicit and explicit einsum mode? Explicit has -> and names the output; implicit lets the library choose (alphabetical).

11. Key takeaways, go deeper, next

Key takeaways

A tensor is an array with any number of axes; the shape lists the size of each.
Shape bugs are the most common bugs; read shapes first.
Einsum: letters in both inputs but missing from the output are summed; letters in the output stay.
Matrix multiplication, dot product, trace, transpose, outer product and attention scores are all one-line einsums.
"Rank" has two meanings; be careful.

Go deeper

Try rewriting a three-loop matrix multiplication as an einsum and compare the answers.
Look up einsum_path in the NumPy documentation and test a chain of three matrices.
Write down the shapes at every step of a small transformer attention block.
Practise broadcasting: add a (3,) bias to a (2, 3) matrix, then to a (3, 2) matrix, and see which fails.

Next on the roadmap: 1.25 Jacobian and Hessian, where we meet matrices of derivatives.

12. Sources and further reading

Accuracy note: the 224 x 224 x 3 selfie, batch of 32 and all small matrices are illustrative examples; the einsum outputs were verified with NumPy, and the NCHW/NHWC defaults are common conventions, not laws.

Search This Blog