Einsum, Tensor Contractions & Convolution
The real notation behind attention, batched ops, and conv-as-matmul.
On this page
Beginner: einsum ("Einstein summation") is a compact notation for describing sums, transposes, and multiplications across the axes of one or more tensors, all in a single short expression — instead of writing nested loops or memorizing which function name does which specific operation.
Intermediate: the pattern is always the same: label each tensor's axes with letters, then say which output axes survive. Repeated letters across inputs get multiplied and summed over ("contracted"); letters that appear in the output are kept. Matrix multiplication is 'ij,jk->ik'; a dot product is 'i,i->'; a batch of matrix multiplications is 'bij,bjk->bik'.
Advanced: once you're comfortable with einsum, the core computation of attention — scores = queries · keysᵀ, batched across many heads and sequence positions at once — is a single readable line instead of a tangle of reshapes and transposes, which is exactly why real Transformer implementations lean on it heavily.
The einsum string is a direct, literal transcription of the summation formula — which is the entire point of the notation.
Chaining three matrices, ABC, is associative (section 1.5) — the result is identical regardless of grouping — but the amount of work is not. Suppose A is (100×5), B is (5×100), and C is (100×5).
Computing (AB)C: first AB costs about 100×5×100 = 50,000 multiplications and produces a (100×100) matrix; then (AB)C costs 100×100×5 = 50,000 more — about 100,000 total.
Computing A(BC) instead: BC costs 5×100×5 = 2,500 and produces a tiny (5×5) matrix; then A(BC) costs 100×5×5 = 2,500 more — only 5,000 total, a 20× saving, purely from choosing a smarter grouping of the exact same product.
Where this is used: this is exactly the optimization opt_einsum (section 1.18's expert note) performs automatically for any chain of three or more tensors — searching for the contraction order with the lowest total cost before executing anything.
Walk through einsum('ij,jk->ik', A, B) by hand for A = [[1, 2], [3, 4]] and B = [[5, 6], [7, 8]]. Output entry (0,0): j is repeated (contracted), so sum over j of A[0,j]·B[j,0] = (1)(5) + (2)(7) = 19. Output entry (0,1): sum over j of A[0,j]·B[j,1] = (1)(6) + (2)(8) = 22. Repeating for every (i,k) pair reproduces exactly the ordinary matrix product A @ B — because that's precisely what this einsum string was written to mean.
That last line is, almost verbatim, the first step of every attention layer in every Transformer — this is what "attention is just dot products" looks like in real code.
#include <cstdio>
#include <vector>
#include <cmath>
#include <random>
using Matrix = std::vector<std::vector<double>>;
using Tensor3 = std::vector<Matrix>;
Matrix randMatrix(int rows, int cols, std::mt19937& rng) {
std::uniform_real_distribution<double> dist(0.0, 1.0);
Matrix M(rows, std::vector<double>(cols));
for (auto& row : M)
for (double& v : row) v = dist(rng);
return M;
}
Tensor3 randTensor3(int batch, int rows, int cols, std::mt19937& rng) {
Tensor3 T;
for (int b = 0; b < batch; ++b) T.push_back(randMatrix(rows, cols, rng));
return T;
}
Matrix matmul(const Matrix& A, const Matrix& B) {
int n = static_cast<int>(A.size());
int k = static_cast<int>(A[0].size());
int m = static_cast<int>(B[0].size());
Matrix C(n, std::vector<double>(m, 0.0));
for (int i = 0; i < n; ++i)
for (int j = 0; j < m; ++j)
for (int x = 0; x < k; ++x) C[i][j] += A[i][x] * B[x][j];
return C;
}
// 'ij,jk->ik' -- j is repeated (contracted), i and k survive: exactly matmul.
Matrix einsumIjJkIk(const Matrix& A, const Matrix& B) { return matmul(A, B); }
// 'bij,bjk->bik' -- one matmul per batch element.
Tensor3 einsumBijBjkBik(const Tensor3& A, const Tensor3& B) {
Tensor3 out;
for (size_t b = 0; b < A.size(); ++b) out.push_back(matmul(A[b], B[b]));
return out;
}
// 'bqd,bkd->bqk' -- simplified scaled dot-product attention scores.
Tensor3 attentionScores(const Tensor3& Q, const Tensor3& K, double scale) {
Tensor3 out;
for (size_t b = 0; b < Q.size(); ++b) {
int seqQ = static_cast<int>(Q[b].size());
int seqK = static_cast<int>(K[b].size());
int dim = static_cast<int>(Q[b][0].size());
Matrix mat(seqQ, std::vector<double>(seqK, 0.0));
for (int q = 0; q < seqQ; ++q)
for (int k = 0; k < seqK; ++k) {
double s = 0.0;
for (int d = 0; d < dim; ++d) s += Q[b][q][d] * K[b][k][d];
mat[q][k] = s / scale;
}
out.push_back(mat);
}
return out;
}
int main() {
std::mt19937 rng(0);
Matrix A = randMatrix(3, 4, rng);
Matrix B = randMatrix(4, 5, rng);
Matrix C1 = matmul(A, B);
Matrix C2 = einsumIjJkIk(A, B);
bool same = true;
for (size_t i = 0; i < C1.size() && same; ++i)
for (size_t j = 0; j < C1[0].size() && same; ++j)
if (std::fabs(C1[i][j] - C2[i][j]) > 1e-9) same = false;
std::printf("%s\n", same ? "true" : "false");
Tensor3 batchA = randTensor3(8, 3, 4, rng);
Tensor3 batchB = randTensor3(8, 4, 5, rng);
Tensor3 batchC = einsumBijBjkBik(batchA, batchB);
std::printf("%zu %zu %zu\n", batchC.size(), batchC[0].size(), batchC[0][0].size());
Tensor3 queries = randTensor3(2, 6, 16, rng);
Tensor3 keys = randTensor3(2, 6, 16, rng);
Tensor3 scores = attentionScores(queries, keys, std::sqrt(16.0));
std::printf("%zu %zu %zu\n", scores.size(), scores[0].size(), scores[0][0].size());
return 0;
}- Convolutional layers can be implemented as a structured matrix multiplication using the "im2col" technique — unfolding overlapping image patches into rows of a matrix, then computing the whole convolution as a single matmul. This is exactly why hardware optimized for matmul (Tensor Cores, section 1.5) is also fast at convolution.
- Toeplitz matrices (constant along each diagonal) are the precise linear-algebra object behind 1D convolution — multiplying by a Toeplitz matrix is convolving with a fixed kernel.
- Modern ML compilers (XLA, TVM, Triton) can fuse and optimize einsum-style expressions automatically, often outperforming manually written loops or even hand-tuned library calls.
- Getting an einsum string subtly wrong (e.g. mismatched or misplaced repeated letters) — it usually still runs, just produces a silently wrong shape or wrong numbers. Always check
.shapeof the result against what you expected. - Forgetting that any letter missing from the output string is summed over — a common mistake is accidentally contracting over an axis you actually wanted to keep.
Going deeper
The name comes from Einstein's own notational shortcut in general relativity and tensor calculus: when the same index appears twice in a term, summation over it is implied without writing a Σ — modern ML libraries adopted the same convention because tensor contractions are exactly as central to deep learning as they are to physics.
At the master level: einsum expressions have a well-defined but non-trivial optimal contractionorder for more than two tensors — a naive left-to-right evaluation can be asymptotically far slower than an optimally ordered one, which is why libraries like opt_einsum exist specifically to search for the cheapest contraction path before executing anything.