Einsum gets slow with multiple arguments

[perimosocordiae]

With multiple arguments, einsum blows up fast:

In [1]: import numpy as np

In [2]: a = np.random.random((5,5))

In [3]: %timeit a.dot(a).dot(a).dot(a)
100000 loops, best of 3: 3.3 µs per loop

In [4]: %timeit np.linalg.multi_dot((a,a,a,a))
10000 loops, best of 3: 30.8 µs per loop

In [5]: %timeit np.einsum('ij,kl,mn,op->ip', a,a,a,a)
100 loops, best of 3: 7.2 ms per loop

In [6]: a = np.random.random((50,50))

In [7]: %timeit a.dot(a).dot(a).dot(a)
10000 loops, best of 3: 82.1 µs per loop

In [8]: %timeit np.linalg.multi_dot((a,a,a,a))
10000 loops, best of 3: 150 µs per loop

In [9]: %timeit np.einsum('ij,kl,mn,op->ip', a,a,a,a)
# takes so long I had to kill IPython.

Is there some fundamental limitation involved, or is this a bug?

[njsmith]

I'm not an einsum expert, so I don't know what "ij,kl,mn,op->ip" means. But it isn't the same as dot:

In [12]: np.max(np.abs(np.einsum("ij,kl,mn,op->ip", a, a, a, a) - a.dot(a).dot(a).dot(a)))
Out[12]: 1629.9863026164651

You want something like "ij,jk,kl,lm":

In [13]: np.max(np.abs(np.einsum("ij,jk,kl,lm", a, a, a, a) - a.dot(a).dot(a).dot(a)))
Out[13]: 7.1054273576010019e-15

This is still slower than dot but by a much less embarrassing margin:

In [15]: %timeit a.dot(a).dot(a).dot(a)
100000 loops, best of 3: 2.58 µs per loop

In [16]: %timeit np.einsum("ij,jk,kl,lm", a, a, a, a)
10000 loops, best of 3: 57.7 µs per loop

[jaimefrio]

Someone more in the know may want to correct me, but I looked at something related to this a while back, and my conclusion was that what tripped einsum was not using any intermediate storage.

A much simpler example:

In [2]: a, b, c =  np.random.rand(3, 1000)

In [3]: %timeit a * np.einsum('i,i', b, c)
100000 loops, best of 3: 6.38 us per loop

In [4]: %timeit np.einsum('j,i,i', a, b, c)
1000 loops, best of 3: 1.16 ms per loop

In [5]: np.allclose(a * np.einsum('i,i', b, c), np.einsum('j,i,i', a, b, c))
Out[5]: True

In the first instance, the dot product of b and c is computed once, and that value is multiplied by every item in a. In the second case, it is redoing that calculation for every item in a.

I don't think this qualifies as a bug, and I am not even sure a general case enhancement is a viable option: you would need to come up with an algorithm that detects groups of operands that could be collapsed, and if some of the operands belong to more than one group, decide in what order to do the collapsing. For the chained matrix multiplication this is solved with a more or less standard DP algorithm, as in multi_dot, for the general case it looks like a daunting task. Although perhaps for 3 or 4 operands a brute force search is a viable option...

[perimosocordiae]

@njsmith: Thanks, that's what I get for not sanity checking. Your example is the one I had initially run across.

@jaimefrio Ok, that makes sense. I think my mental model for how einsum works was off. I'm not sure how to make the docs clearer, though.

[charris]

Closing.