Einsum and hypergraph categories

einsum is a small language for combining tensors, made popular by numpy. This page explains what it does with matrices, then shows that an einsum specification is a picture in a hypergraph category, so that the same specification means something in any such category: matrices, relations, or string diagrams. The last part uses the pictures to see why the order in which tensors are combined matters, and how Proarrow.Tools.Einsum chooses it.

Einsum on matrices

Einstein's summation convention writes the product of two matrices as

$$C_{ik} = \sum_j A_{ij} B_{jk}$$

and then leaves out the sum: an index that occurs twice on the right and not on the left is summed. einsum makes the convention a string. The product above is "ij,jk->ik": the letters before the arrow are the indices of the inputs, one group per tensor, and the letters after it are the indices of the result. The rules are:

• each letter is an index, and every occurrence of a letter takes the same value;
• the entries of the inputs at those values are multiplied;
• a letter that is not in the output is summed over.

In proarrow:

einsum @"ij,jk->ik" a b

A tensor is a Tensor '[i, j], whose type lists the objects of its indices, so in the category of matrices Mat the type also lists the dimensions. The type of the result follows from the types of the inputs (a letter used with two different dimensions is a type error):

EinsumType "ij,jk->ik" '[ '[M Nat2, M Nat3], '[M Nat3, M Nat4]]
  = Tensor '[M Nat2, M Nat3] -> Tensor '[M Nat3, M Nat4] -> Tensor '[M Nat2, M Nat4]

The standard examples, with a matrix A, a matrix B and vectors a and b:

specification	result	name
"ij,jk->ik"	$\sum_j A_{ij} B_{jk}$	matrix product
"ij->ji"	$A_{ji}$	transpose
"ii->"	$\sum_i A_{ii}$	trace
"ii->i"	$A_{ii}$	diagonal
"ij->"	$\sum_{i,j} A_{ij}$	sum of the entries
"ij->i"	$\sum_j A_{ij}$	row sums
"i,i->"	$\sum_i a_i b_i$	dot product
"i,j->ij"	$a_i b_j$	outer product
"ij,ij->ij"	$A_{ij} B_{ij}$	entrywise product
"i->ii"	$a_i$ on the diagonal, 0 elsewhere	diagonal matrix

The last one is not allowed in numpy, which forbids a letter twice in the output. Here both positions of the output are the same index i, so an entry where the two positions differ has no value of i to take, and is 0: the entry at i, j is $a_i$ when $i = j$ and 0 otherwise. Without an arrow, the output is, as in numpy, the letters used once, in alphabetical order: "ij,jk" is "ij,jk->ik".

Einsum is a picture

Draw each tensor as a box with one wire for each of its indices, and connect all the wires that have the same letter. A letter in the output also gets a wire to the bottom of the picture. For the matrix product, A has wires i and j, B has j and k, the two j wires are joined, and i and k go to the output.

Joining wires is not just a matter of drawing. A letter can occur on any number of wires, two in the matrix product, three in "ij,ij->ij" (twice in the inputs, once in the output), one in "ij->". So the picture is not a graph, where an edge has two ends, but a hypergraph: each letter is a node that any number of wire ends can meet in. Together with the output letters, which form its boundary, an einsum specification is an open hypergraph.

What does a node with $n$ wire ends coming in and $m$ going out mean? In matrices it is the tensor that is 1 when all $n + m$ indices are equal and 0 otherwise: a node in the middle of the picture forces its wires to agree, and as it has no wire of its own to the outside it is summed over. These tensors are called spiders, and they are made from four basic ones:

• merge, from two wires to one, which is 1 when the two indices and the result are all equal;
• unit, from no wires to one, the vector of ones;
• copy, from one wire to two, which is 1 when the index and the two results are all equal;
• discard, from one wire to none, which sums.

Merge and unit form a commutative monoid, copy and discard a cocommutative comonoid, and together they form a special commutative Frobenius algebra. The spider theorem says that any connected network built from these four is equal to a single spider, the one with the same number of wires in and out. So it does not matter how the wires of a letter are joined: every way of joining them means the same.

A hypergraph category is a symmetric monoidal category where every object has such a Frobenius algebra, compatible with the tensor. Fong and Spivak's Hypergraph categories is the standard reference, and Bonchi, Gadducci, Kissinger, Sobociński and Zanasi show in String diagram rewrite theory I (Theorem 4.1) that open hypergraphs are the morphisms of the free hypergraph category on a set of boxes. Einsum is a syntax for those morphisms: a specification is an open hypergraph, so it has a meaning in every hypergraph category. In proarrow einsum works in any of them:

• in Mat, the matrices, it is numpy's einsum, with the arithmetic of the entries;
• in FinRel, the relations between finite sets, merge is equality and discard is "there is some", so "ij,jk->ik" is composition of relations and an einsum in general is a conjunctive query, a natural join;
• in SVG and DOT it draws the picture;
• in TensorNetwork, the same matrices as in Mat, kept as a network of tensors instead of as one matrix, so that it is fast (see the last section);
• in the open hypergraphs themselves (Proarrow.Category.Instance.OpenHypergraph), the free hypergraph category, it gives the hypergraph of the specification, so two specifications are equal there when and only when they are equal in every hypergraph category.

To compute an einsum, einsum builds the open hypergraph of its specification and then reads it back into the category it is used in. All the pictures below are that read-back into SVG.

How the features are drawn

A tensor is a box without inputs. A point on a wire is one of the four basic spiders: hollow for merge and unit, filled for copy and discard. A merge followed by a discard, a spider from two wires to none, is drawn as a cap, the bend that joins two wires, and the points are drawn next to what they meet.

specification	picture	what to look for
"ij,jk->ik"		the summed j is a cap between A and B
"ij->ji"		a crossing
"ii->"		a cap between two wires of A
"ii->i"		a merge of two wires of A
"ij->"		two discards
"ij->i"		one discard
"i,i->"		a cap between a and b
"i,j->ij"		nothing joined
"ij,ij->ij"		two merges: two wires in, one out
"i->ii"		a copy

Everything at once

Here is a specification that uses all of these at once:

einsum @"kl,jkk,ijx,li->li" v u t w

• The four tensors form a cycle: V and U share k, U and T share j, T and W share i, W and V share l.
• U has k twice, so only its diagonal is used.
• x occurs only in T and not in the output, so T is summed over x before it meets anything.
• i and l each join two tensors and also go to the output, so they are spiders with three legs.
• The output lists l before i.

The simplest way to compute it is the one the rules suggest: multiply everything, then join up the wires. In the picture that is all four tensors side by side, followed by one spider for each letter.

Take the tensors to be over the dimensions 2 for i and k and 3 for j, x and l. Then V has 6 entries, U 12, T 18 and W 6, and the tensor product of all four, before anything is summed, has 6 · 12 · 18 · 6 = 7776 entries. The result has 6.

Two tensors at a time

Every way of drawing the picture means the same, by the spider theorem, so we are free to choose a cheaper one. Three changes give one:

• Contract the tensors two at a time: put two pieces that share a letter side by side, join them, and keep going until one piece is left.
• Join a letter as soon as nothing still to come has it: its spider then goes right after the join that brings in its last tensor, and if the letter is not in the output its wire ends there.
• Before a tensor meets the rest, apply the part of the spiders that only concern it: merge its own wires of one letter, as U's two ks, and discard the letters nothing else has, as T's x.

What is left is the order. A first choice is a greedy one that only counts wires: next is the pair after whose join the fewest wires are open. In this example every pair leaves two wires open, and the first one wins: V and U, which share k, then T joins the result, then W:

The picture now shows the order of evaluation as a tree: each join is two pieces side by side, followed by the spiders that become possible. With the dimensions above, the largest product along the way is the piece of V and U, with l and j, times what is left of T, i and j: 9 · 6 = 54 entries.

The gain is largest for long chains: for four 3 × 3 matrices, "ij,jk,kl,lm->im", multiplying everything first builds a tensor of 6561 entries, and two at a time never more than 81. In Mat that took the computation from over ten minutes to a few hundredths of a second.

Counting sizes

Counting wires treats every index as equally big. Numpy's einsum chooses its order with sizes: it contracts first the pair of tensors for which the size of the result, less the sizes of the two tensors, is smallest, and then keeps going greedily. Proarrow.Tools.Einsum does the same, with the size of each object given by the class Sized: the dimension in Mat, the number of elements in FinRel.

For this example that starts with U and T. After U merges its two ks and T discards x, each has 6 entries, and joining them over j leaves k and i, 4 entries: the best pair. Then V joins the result, and then W:

Now the largest product along the way is U and T side by side, 6 · 6 = 36 entries, where multiplying everything first needed 7776.

A category made for tensors

A better order is not the whole story. In Mat every arrow is one matrix, and that includes the arrows that only rearrange wires: a swap, a copy or a merge, an identity next to a tensor. Each step of the picture is then a full matrix of the size of all the wires it passes, so the four-matrix chain at dimension 6 still takes over two minutes there, although the contractions it needs are tiny.

Proarrow.Category.Instance.TensorNetwork computes the same matrices, but keeps an arrow as the picture itself. An object is a list of dimensions, one for each wire. An arrow has nodes, each with a dimension, a node for each of its input and output wires, and a list of dense tensors, each with a node for each of its indices. Its entries are what the spider rules say: sum over the indices of the nodes that agree with those of the wires, and multiply the entries of the tensors.

Then:

• swaps, copying, merging, discarding, cups and caps are wirings without tensors, so they cost nothing;
• putting two arrows side by side puts their tensors side by side, without multiplying them;
• composing glues the two wirings and sums out the nodes that are on no wire any more, multiplying only the tensors such a node joins. That is one contraction of the kind numpy's einsum is made of, so the order chosen in the previous sections is the order in which the work is done.

The four-matrix chain at dimension 256 takes about a third of a second in TensorNetwork.

In the theory this is a category of decorated corelations, after Brendan Fong's Decorated corelations: the wiring is a cospan in which every node is on a wire, and the tensors are the decoration. Composing glues the cospans and keeps only the nodes still on a wire. A node dropped in this way is summed out of the tensors, and a closed loop, which corelations alone would forget, leaves its dimension behind as a scalar.