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