Einstein Summation: einsum

Reviewed & published by Brayan K

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.

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

• Read the einsum 'inputs->output' subscript notation
• Dot product with 'i,i->'
• Matrix multiply with 'ij,jk->ik'
• Transpose with 'ij->ji'
• Row and column sums with 'ij->i' and 'ij->j'
• Trace of a matrix with 'ii->'

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.

import numpy as np

a = np.array([1, 2, 3])
b = np.array([4, 5, 6])

np.einsum("i,i->", a, b)    # 1*4 + 2*5 + 3*6 = 32
# i appears in inputs but not output -> summed away

import numpy as np

a = np.array([1, 2, 3])
b = np.array([4, 5, 6])

# Dot product: multiply matching elements, sum them
print(np.einsum("i,i->", a, b))   # 32
print(np.dot(a, b))               # 32  (same thing)

# Element-wise product, keeping the i axis -> NOT summed
print(np.einsum("i,i->i", a, b))  # [ 4 10 18]

# ✅ Expected output:
# 32
# 32
# [ 4 10 18]

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.

A = np.array([[1, 2],
              [3, 4]])
B = np.array([[5, 6],
              [7, 8]])

np.einsum("ij,jk->ik", A, B)   # matrix product, same as A @ B
np.einsum("ij->ji", A)         # transpose
np.einsum("ij->i", A)          # row sums  -> [3 7]
np.einsum("ij->j", A)          # column sums -> [4 6]

import numpy as np

A = np.array([[1, 2],
              [3, 4]])
B = np.array([[5, 6],
              [7, 8]])

# Matrix product (sum over shared j)
print(np.einsum("ij,jk->ik", A, B))
# [[19 22]
#  [43 50]]
print(np.array_equal(np.einsum("ij,jk->ik", A, B), A @ B))   # True

# Transpose by reordering output letters
print(np.einsum("ij->ji", A))
# [[1 3]
#  [2 4]]

# Row sums and column sums
print(np.einsum("ij->i", A))   # [3 7]
print(np.einsum("ij->j", A))   # [4 6]

# ✅ Expected output:
# [[19 22]
#  [43 50]]
# True
# [[1 3]
#  [2 4]]
# [3 7]
# [4 6]

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.

M = np.array([[1, 2],
              [3, 4]])

np.einsum("ii->i", M)    # diagonal -> [1 4]
np.einsum("ii->", M)     # trace    -> 5  (1 + 4)

import numpy as np

M = np.array([[1, 2],
              [3, 4]])

# Diagonal as a vector
print(np.einsum("ii->i", M))   # [1 4]

# Trace = sum of the diagonal
print(np.einsum("ii->", M))    # 5
print(np.trace(M))             # 5  (same)

# Outer product: no shared index, keep both
u = np.array([1, 2, 3])
v = np.array([10, 20])
print(np.einsum("i,j->ij", u, v))
# [[10 20]
#  [20 40]
#  [30 60]]

# ✅ Expected output:
# [1 4]
# 5
# 5
# [[10 20]
#  [20 40]
#  [30 60]]

🎯 YOUR TURN: Fill in the Blank

Complete the subscript string so this einsum performs a matrix multiplication of A and B.

import numpy as np

A = np.array([[1, 0],
              [0, 1]])      # identity
B = np.array([[5, 6],
              [7, 8]])

# Fill in the subscripts for matrix multiplication: keep i and k, sum j
result = np.einsum("___", A, B)

print(result)
# Expected output:
# [[5 6]
#  [7 8]]

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)

❌ Mismatched index letters between inputs

Using a fresh letter for the shared axis breaks the contraction:

np.einsum("ij,kl->ik", A, B)   # ❌ j and l never meet, this is not matmul

✅ Fix: reuse one letter for the axis you want to sum:

np.einsum("ij,jk->ik", A, B)   # ✅ shared j is contracted

❌ 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.

import numpy as np

# Rows of scores, columns are categories
scores = np.array([[2, 4, 1],
                   [5, 0, 3]])

weights = np.array([10, 1, 100])   # weight per category

# For each row: sum over j of scores[i, j] * weights[j]
totals = np.einsum("ij,j->i", scores, weights)

print(totals)
# Row 0: 2*10 + 4*1 + 1*100 = 124
# Row 1: 5*10 + 0*1 + 3*100 = 350
# [124 350]

# Cross-check with a matrix-vector product
print(np.array_equal(totals, scores @ weights))   # True

# ✅ Expected output:
# [124 350]
# True

❓ 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.

Practice quiz

Continue this course

• Previous: Vectorizing Python Functions
• Next: Checkpoint: Numerical Computing — Combine grids, polynomial fits and einsum in a build — then a quiz
• Quick reference: NumPy cheat sheet

Frequently asked questions

How do I read an einsum subscript string?

The string has the form 'inputs->output'. Each letter labels an axis. Letters that appear in the inputs but NOT in the output are summed over; letters that appear in both are kept. So 'ij,jk->ik' means multiply along the shared j axis and keep i and k, which is exactly matrix multiplication. 'ii->' keeps no axes and sums the diagonal, giving the trace.

Why use np.einsum instead of np.dot or @ ?

For a plain matrix product, @ or np.dot is clearer and you should use them. einsum shines when an operation is awkward to express otherwise: batched products, contracting specific axes of higher-dimensional tensors, simultaneous transposes and sums, or operations that would otherwise need several reshape and transpose steps. It expresses all of them in one readable line.

Does a repeated index always mean summation?

A repeated index is summed over only when it does not appear in the output. In 'ij,jk->ik' the index j is repeated across the two inputs and is absent from the output, so it is summed. In 'ii->i' the repeated i appears in the output, so instead of summing it selects the diagonal. The output side of the arrow controls what is kept versus contracted.

Can einsum transpose and sum at the same time?

Yes. Because the output subscripts can list axes in any order, 'ij->ji' transposes. You can also drop an axis to sum it: 'ij->i' sums each row and 'ij->j' sums each column. Combining reorder and contraction in a single call is one of einsum's biggest strengths.