einsum¶

einsum writes a contraction, the thing that matrix products, dot products, traces and transposes all are, in one notation. It is Einstein's summation convention with the sum signs left out: a letter that appears on the input side and not on the result side is summed over. The matrix product

\[ C_{ik} = \sum_j A_{ij}\, B_{jk} \]

is einsum<"ij,jk->ik">(a, b). The operands are listed before the arrow, each with one letter per axis; the letters after the arrow name the axes of the result. j is in both operands and not in the result, so it is summed.

It is NumPy's einsum in its explicit form (the one with an arrow), and the result is lazy like any expression.

The notation¶

Operands are separated by ,, then comes -> and the result. The arrow is required.
A label is one letter, a to z, and stands for one axis: an operand of rank 2 has two letters. No spaces, no capitals, no ....
The result lists labels that appear in some operand, each at most once. It can be empty, and then the result has rank 0 and is read as result[].
A label that is not in the result is summed over.
A label that repeats within one operand takes its diagonal. A label in several operands must have the same extent in each.

The result has one axis per label of the result, in the order written. Its extent is that of the first operand axis carrying the label, static where the operand's type fixes it.

Some notations¶

Notation	Formula	What it is
`"ij,jk->ik"`	\(C_{ik} = \sum_j A_{ij} B_{jk}\)	matrix product
`"ij->ji"`	\(B_{ji} = A_{ij}\)	transpose
`"ij->i"`	\(r_i = \sum_j A_{ij}\)	row sums
`"ii->i"`	\(d_i = A_{ii}\)	diagonal
`"ii->"`	\(\operatorname{tr} A = \sum_i A_{ii}\)	trace
`"i,i->"`	\(u \cdot v = \sum_i u_i v_i\)	dot product
`"i,j->ij"`	\(C_{ij} = u_i v_j\)	outer product
`"ij,ij->"`	\(\sum_{ij} A_{ij} B_{ij}\)	Frobenius inner product
`"ijk,i,j->k"`	\((a \times b)_k = \sum_{ij} \varepsilon_{ijk}\, a_i b_j\)	cross product, three operands

In code¶

Every block on this page is part of examples/10_einsum.cpp, which is built and run with the tests. The output shown is its output. The operands are two 2 × 2 matrices,

\[ A = \begin{pmatrix} 1 & 2 \\ 3 & 4 \end{pmatrix}, \qquad B = \begin{pmatrix} 5 & 6 \\ 7 & 8 \end{pmatrix}, \]

and two vectors of three numbers, \(u = (1, 2, 3)\) and \(v = (4, 5, 6)\):

using strata::einsum;
using strata::Shape;
using strata::slice;
using strata::Tensor;
using strata::TensorView;

using Matrix = TensorView<const double, Shape<2, 2>>;
using Vector = TensorView<const double, Shape<3>>;

// Column-major, as strata's default: [[1, 2], [3, 4]] and [[5, 6], [7, 8]].
constexpr std::array a_values{1.0, 3.0, 2.0, 4.0};
constexpr std::array b_values{5.0, 7.0, 6.0, 8.0};

const Matrix a{a_values.data(), Shape<2, 2>{}};
const Matrix b{b_values.data(), Shape<2, 2>{}};

One product, one operand¶

// The matrix product: j is in both operands and not in the result, so it is summed.
std::println("a b:\n{}", Tensor{einsum<"ij,jk->ik">(a, b)});   // [[19, 22], [43, 50]]

a b:
[[19, 22],
 [43, 50]]

Element by element, \(C_{00} = 1 \cdot 5 + 2 \cdot 7 = 19\). With one operand, einsum transposes, sums and takes diagonals:

// One operand: a transpose, row sums, the diagonal, the trace.
std::println("a transposed:\n{}", Tensor{einsum<"ij->ji">(a)});
std::println("row sums of a: {}", Tensor{einsum<"ij->i">(a)});   // [3, 7]
std::println("diagonal of a: {}", Tensor{einsum<"ii->i">(a)});   // [1, 4]
std::println("trace of a: {}", einsum<"ii->">(a)[]);             // 5

a transposed:
[[1, 3],
 [2, 4]]
row sums of a: [3, 7]
diagonal of a: [1, 4]
trace of a: 5

The trace has no letters left after the arrow: the result is a number of rank 0, and [] reads it.

Two vectors, three operands¶

// Two vectors: the dot product, and the outer one, of higher rank than either.
constexpr std::array u_values{1.0, 2.0, 3.0};
constexpr std::array v_values{4.0, 5.0, 6.0};
const Vector u{u_values.data(), Shape<3>{}};
const Vector v{v_values.data(), Shape<3>{}};
std::println("u . v: {}", einsum<"i,i->">(u, v)[]);   // 32
std::println("u v^T:\n{}", Tensor{einsum<"i,j->ij">(u, v)});

u . v: 32
u v^T:
[[ 4,  5,  6],
 [ 8, 10, 12],
 [12, 15, 18]]

The dot product sums over i; the outer product sums over nothing and has a higher rank than either operand. A third operand is one more term in the product. The cross product is the contraction of the Levi-Civita symbol \(\varepsilon_{ijk}\) with both vectors:

\[ (u \times v)_k = \sum_{i,j} \varepsilon_{ijk}\, u_i\, v_j \]

// Three operands: the cross product through the Levi-Civita symbol, (a x b)_k = eps_ijk a_i b_j.
Tensor<double, Shape<3, 3, 3>> eps{};
eps[0, 1, 2] = eps[1, 2, 0] = eps[2, 0, 1] = 1.0;
eps[0, 2, 1] = eps[2, 1, 0] = eps[1, 0, 2] = -1.0;
std::println("u x v: {}", Tensor{einsum<"ijk,i,j->k">(eps, u, v)});   // [-3, 6, -3]

u x v: [-3,  6, -3]

Lazy, and an expression as an operand¶

einsum computes nothing when it is called. The result is an expression, and reading one element computes one sum. So a single element, or a single row, costs only what it needs:

// Lazy: an element, or a row, computes only what it needs. Row 1 of a b is row 1 of a times b.
const auto product = einsum<"ij,jk->ik">(a, b);
std::println("(a b)[1, 0]: {}, row 1: {}", product[1, 0], Tensor{product | slice<0>(1)});

(a b)[1, 0]: 43, row 1: [43, 50]

Any expression is an operand, and an einsum is one too. \((A + A)\,B\), and \((AB)^\mathsf{T} = B^\mathsf{T} A^\mathsf{T}\) as two transposes and a product:

// Any expression is an operand, and an einsum is one too: (a + a) b, and (a b)^T = b^T a^T.
std::println("(a + a) b:\n{}", Tensor{einsum<"ij,jk->ik">(a + a, b)});
std::println("b^T a^T:\n{}", Tensor{einsum<"ij,jk->ik">(einsum<"ij->ji">(b), einsum<"ij->ji">(a))});

(a + a) b:
[[ 38,  44],
 [ 86, 100]]
b^T a^T:
[[19, 43],
 [22, 50]]

What it costs¶

einsum is one loop over every letter at once, with no intermediate tensors. For one product of two matrices that is the textbook cost, \(n_i n_j n_k\) multiplications.

For a chain it is not the best way. Contracting \(A_{ij} B_{jk} C_{kl}\) in a single einsum runs a loop over all four letters,

\[ n_i\, n_j\, n_k\, n_l \]

multiplications, where two products in turn, \((AB)\) and then \((AB)\,C\), cost only

\[ n_i\, n_j\, n_k + n_i\, n_k\, n_l . \]

For square matrices of size \(n\) that is \(n^4\) against \(2n^3\). So a chain of three or more matrices is cheaper as "ij,jk->ik" twice than as one einsum.

An operand that is itself an expression is computed again for each value of the letters it lacks: above, \(A + A\) once for every k. Where that is dear, the assignment computes it into a buffer first and reads the buffer (the cost model decides, at compile time). A tensor of its values, made by hand, does the same:

const Tensor twice_a = a + a;
std::println("(a + a) b, made first:\n{}", Tensor{einsum<"ij,jk->ik">(twice_a, b)});

(a + a) b, made first:
[[ 38,  44],
 [ 86, 100]]

What is rejected¶

A notation that is wrong, or operands that do not fit it, is a compile error, and the message begins with einsum:. That covers:

no ->, or a character other than letters, , and the arrow;
a label repeated in the result, or in the result and in no operand;
a different number of operands than the text has, or an operand whose rank is not the number of its labels;
operands whose types fix different extents for one label.

Extents that only dynamic extents fix, and that differ, are caught at run time, as a precondition violation.