# 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

## 1 Answer

It seems that the general problem of finding the optimal contraction order is NP-hard [1]. A recent paper on approximately optimizing the contraction order, and containing relevant references, is [2].

[1] Chi-Chung, Lam, P. Sadayappan, and Rephael Wenger. "On optimizing a class of multi-dimensional loops with reduction for parallel execution." Parallel Processing Letters 7.02 (1997): 157-168.

[2] Schindler, Frank, and Adam Jermyn. "Algorithms for tensor network contraction ordering." Machine Learning: Science and Technology (2020).

• Thank you, "tensor network contraction" is the keyword I was missing – Arthur B Aug 20 at 10:30
• 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
• Sorry, yes I meant to write quartic. And the answer to my question is that, yes, this is a thing, but there's a flag... optimized-einsum.readthedocs.io/en/stable/index.html – Arthur B Aug 20 at 17:39
• Specifically I was thinking: the paper assumes that contractions always dominate naive evaluation, but surely if that were the case, numpy's einsum would implement at least a naive contraction. Turn out it does, but only with a flag. – Arthur B Aug 20 at 17:41