Einstein Summation: einsum
np.einsum uses a compact subscript notation to describe products and sums over array axes, unifying dot products, matrix multiplication, transposes, traces, and axis sums in a single expressive function.
Learn Einstein Summation: einsum in our free NumPy course — a beginner-friendly interactive lesson with worked examples, a practice exercise and a quick…
Part of the free Numpy course at LearnCodingFast — hands-on lessons with examples you run in your browser, plus practice exercises and a quick quiz.
You'll learn to read the 'inputs->output' notation, compute dot products and matrix products, transpose and sum along axes, and take a trace — all by choosing which index letters to keep.
What You'll Learn in This Lesson
1 Reading the Notation: Dot Products
Every einsum call is a string 'inputs->output' . Each letter names an axis. The rule is simple: any letter that appears in the inputs but not in the output is summed over. For a dot product of two vectors, label both with i and keep nothing: 'i,i->' multiplies element by element and sums the result.
2 Matrix Multiply, Transpose, and Sums
Matrix multiplication shares the inner axis: 'ij,jk->ik' sums over j and keeps i and k . Reordering the output letters transposes: 'ij->ji' . Dropping a letter sums along it, so 'ij->i' gives row sums and 'ij->j' gives column sums.
3 Trace, Diagonal, and the Big Picture
Repeating an index within one input picks out the diagonal. 'ii->i' extracts the diagonal as a vector, and 'ii->' sums it to give the trace . With these few patterns you can see how einsum unifies many separate NumPy functions under one consistent rule about which letters survive to the output.
🎯 YOUR TURN: Fill in the Blank
Complete the subscript string so this einsum performs a matrix multiplication of A and B.
Answer: ij,jk->ik . The shared j is summed; i and k are kept. Multiplying by the identity returns B unchanged.
! Common Errors (And How to Fix Them)
Using a fresh letter for the shared axis breaks the contraction:
✅ Fix: reuse one letter for the axis you want to sum:
❌ Forgetting that missing output letters are summed
'ij->' sums the entire matrix to one number, which may not be what you wanted.
✅ Fix: list every axis you want to keep on the output side, e.g. 'ij->ij' to keep the matrix intact.
❌ Subscript count not matching array dimensions
A 2D array needs two letters; giving it one or three raises a ValueError.
✅ Fix: use exactly one letter per axis, so a shape (n, m) array gets two letters like ij .
🎯 Mini Challenge: Weighted Row Totals
Use a single einsum to multiply each row of a matrix by a weight vector and sum across the columns.
❓ Frequently Asked Questions
Lesson complete — one notation to rule them all!
You can read 'inputs->output' subscripts, and use np.einsum for dot products, matrix multiplication, transposes, axis sums, diagonals, and traces — all by choosing which letters survive.
🚀 Up next: Checkpoint — Numerical Computing — recap everything and tackle a multi-step build challenge.
Practice quiz
In an einsum string, what happens to a letter that appears in the inputs but not the output?
- It is summed over
- It is transposed
- It raises an error
- It is kept unchanged
Answer: It is summed over. Letters missing from the output side are contracted (summed away).
Which subscript string computes a vector dot product of a and b?
- 'i,i->i'
- 'ij,jk->ik'
- 'i,i->'
- 'i->'
Answer: 'i,i->'. 'i,i->' multiplies matching elements and sums them into a scalar.
What does np.einsum('ij,jk->ik', A, B) compute?
- The transpose of A
- Matrix multiplication of A and B
- The trace of A
- An element-wise product
Answer: Matrix multiplication of A and B. The shared j is summed and i, k are kept, which is matrix multiplication.
What does np.einsum('ij->ji', A) do?
- Sums the rows
- Takes the trace
- Sums the columns
- Transposes A
Answer: Transposes A. Reordering the output letters swaps the axes, giving the transpose.
What does np.einsum('ij->i', A) return for A = [[1,2],[3,4]]?
Dropping j sums each row: 1+2=3 and 3+4=7.
What does np.einsum('ii->', M) compute for a square matrix M?
- The diagonal as a vector
- The transpose
- The trace (sum of the diagonal)
- The determinant
Answer: The trace (sum of the diagonal). Repeating i and dropping it from the output sums the diagonal: the trace.
What does np.einsum('ii->i', M) return?
- The trace as a scalar
- The diagonal as a vector
- The first row
- The first column
Answer: The diagonal as a vector. Keeping the repeated index in the output selects the diagonal as a 1D array.
When is np.dot or @ preferable to einsum?
- Never, einsum is always better
- Only for 1D arrays
- For a plain matrix product where @ is clearer
- Only for integer arrays
Answer: For a plain matrix product where @ is clearer. For an ordinary matrix product, @ or np.dot reads more clearly.
What does np.einsum('i,i->', [1,2,3], [4,5,6]) return?
- 32
Answer: 32. 1*4 + 2*5 + 3*6 = 32, the dot product.
How many index letters does a 2D array need in einsum?
- 1
- 2
- 3
- 0
Answer: 2. Exactly one letter per axis, so a 2D array uses two letters like ij.
Continue this course
- Previous: Vectorizing Python Functions
- Next: Checkpoint: Numerical Computing