Subscripts instead of transposes

einsum describes a tensor calculation by naming its axes. For example, this expression computes the dot product of every query and key in each batch item:

q: (batch, queries, features)
k: (batch, keys, features)

torch.einsum("bqd,bkd->bqk", q, k)
result: (batch, queries, keys)

Read the equation

The labels on each input correspond to its axes in order. In this example:

• b identifies the same batch item in both inputs and the output.
• q and k keep query and key positions separate.
• d identifies matching features. It is absent from the output, so those products are summed over the feature axis.

The letters are labels you choose, not reserved names. The right side of -> specifies the output axes and their order. You do not need to transpose k yourself because the equation pairs the two feature axes explicitly.

import torch

q = torch.tensor([[[1.0, 0.0], [0.0, 2.0]]])
k = torch.tensor([[[1.0, 1.0], [2.0, 0.0], [0.0, 3.0]]])
out = torch.einsum("bqd,bkd->bqk", q, k)
print(out.tolist())  # [[[1.0, 2.0, 0.0], [2.0, 0.0, 6.0]]]

This is a dot-product score calculation without attention scaling or softmax. For these inputs it is equivalent to torch.bmm(q, k.transpose(1, 2)).

Other uses of the notation

An axis can be summed even if its label occurs only once: ij->i sums matrix rows. Repeating a label within one input selects a diagonal: ii->i returns a square matrix's diagonal. Neither case requires multiplying two tensors. Explicit output labels make these differences easier to see.

Choosing an implementation

Use the expression that makes the calculation clear, then benchmark if it is performance-sensitive. Eager einsum may perform several internal view or permutation operations around a matrix multiplication. An operation count alone does not tell you elapsed time or GPU kernel count.

Performance depends on shapes, memory layout, device and execution mode. Compilation can change the implementation, but does not guarantee that all overhead disappears or that einsum is faster than a direct matrix product.

Multiply rows by columns

For A with shape (M, K) and B with shape (K, N), the matrix product has shape (M, N). Each output entry is a dot product of one row of A and one column of B:

$$C_{ij}=\sum_{k=1}^{K}A_{ik}B_{kj}.$$

For example, multiplying row [1, 2] by column [3, 4] gives 1 * 3 + 2 * 4 = 11. Elementwise multiplication would instead give [3, 8]; it does not include the sum.

If examples are rows of X, with shape (B, D), and weights have shape (D, H), then XW produces H features for each of the B examples.

Keep batch axes separate from matrix axes

For batched matrix multiplication, the last two axes hold each matrix. Leading axes identify batches of matrices. A general matrix-multiply operation may broadcast compatible batch axes. A stricter batched operation may require them to match exactly. Use the operation specified by the exercise.

Einstein notation, often used with einsum, names the axes explicitly:

• bij,bjk->bik multiplies corresponding matrices in batch b and sums over j.
• bi,bj->bij forms an outer product for each example: output (b, i, j) is the product of input (b, i) and input (b, j).

An axis is summed out when its label is absent from the explicit output. A label that appears in both inputs is not necessarily summed: b stays in both outputs above.

Choose a multiplication order with small intermediate results

In exact arithmetic, (AB)C = A(BC), but the two orders can have very different costs. If the shapes are (1000, 10), (10, 10), and (10, 1), forming BC first needs a (10, 1) intermediate. Forming AB first needs a (1000, 10) intermediate and many more operations.

Floating-point rounding can also differ between orders, so do not expect every bit of the two results to match.