# Compiling einstein sums optimally

Einstein summation is a convenient way to express tensor operations which has found its way in tensor libraries like numpy, torch, tensorflow, etc.

Its flexibility lets us represent the product of three matrices, $$X$$, $$Y$$, $$Z$$ of dimensions $$(a,b)$$, $$(b,c)$$, $$(c,d)$$ as

X.Y.Z = einsum('ab,bc,cd->ad',X,Y,Z)


However, the above compiles to something like

for a_ in range(a):
for d_ in range(d):
res[a_,d_] = 0
for b_ in range(b):
for c_ in range(c):
res[a_,d_] += X[a_,b_] * Y[b_,c_] * Z[c_, d_]


This native is quadratic in the size of the matrices when simply doing

einsum('ac,cd->ad',einsum('ab,bc'->'ac', X, Y), Z)


Would be merely cubic.

There are roughly three levels of optimization we can imagine a smarter implementation of einsum to perform.

1. Decompose an einsum of tensors $$(x_1, \ldots, x_n)$$ into an einsum of pairs of tensors $$x_1, x_2$$, $$e(x_1, x_2), x_3$$, etc to optimize computation time.

2. Rely on associativity (where applicable) to pick those pairs judiciously (this is a classic dynamic programming problem) and construct the appropriate intermediary tensors.

3. Discover Strassen-like formulas for the specific tensor computation

While 3 seems clearly out of reach, 1 and 2 seem like they could be achievable exactly with a reasonably straightforward algorithm. Are such algorithms known for generic einstein summations? Have they been studied?

• What does the notation cd->ab mean? – Mark Aug 18 at 15:53
• The quote inserted was a typo, it's ab,cd->ad – Arthur B Aug 18 at 21:36

• Not all evaluation strategies are contractions though. The naive product of 3 $n \times n$ matrices in $\mathcal{O}(n^4)$ time doesn't happen if you follow any one of the two contraction strategy... yet, you'll see numpy's einsum for instance run it in quadratic time anyway. What gives? – Arthur B Aug 20 at 10:38
• By quadratic, do you actually mean quartic (${\cal O}(n^4)$)? I guess that any optimal evaluation strategies would consist of contractions. – smapers Aug 20 at 10:49