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?

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

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).

  • 1
    $\begingroup$ Thank you, "tensor network contraction" is the keyword I was missing $\endgroup$ – Arthur B Aug 20 '20 at 10:30
  • $\begingroup$ 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? $\endgroup$ – Arthur B Aug 20 '20 at 10:38
  • $\begingroup$ By quadratic, do you actually mean quartic (${\cal O}(n^4)$)? I guess that any optimal evaluation strategies would consist of contractions. $\endgroup$ – smapers Aug 20 '20 at 10:49
  • $\begingroup$ 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 $\endgroup$ – Arthur B Aug 20 '20 at 17:39
  • $\begingroup$ 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. $\endgroup$ – Arthur B Aug 20 '20 at 17:41

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.