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.

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.

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.

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.

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.

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 .

Use a single einsum to multiply each row of a matrix by a weight vector and sum across the columns.

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.