8

I am trying to understand the einsum function in NumPy. In this documentation, the last example, 

>>> a = np.arange(60.).reshape(3,4,5)
>>> b = np.arange(24.).reshape(4,3,2)
>>> np.einsum('ijk,jil->kl', a, b)
array([[ 4400.,  4730.],
       [ 4532.,  4874.],
       [ 4664.,  5018.],
       [ 4796.,  5162.],
       [ 4928.,  5306.]])
>>> np.einsum(a, [0,1,2], b, [1,0,3], [2,3])
array([[ 4400.,  4730.],
       [ 4532.,  4874.],
       [ 4664.,  5018.],
       [ 4796.,  5162.],
       [ 4928.,  5306.]])
>>> np.tensordot(a,b, axes=([1,0],[0,1]))
array([[ 4400.,  4730.],
       [ 4532.,  4874.],
       [ 4664.,  5018.],
       [ 4796.,  5162.],
       [ 4928.,  5306.]])

I don't understand what's going on with this np.einsum('ijk,jil->kl', a, b) function. Can someone express it in a more explicit way, something like $$\sum_{???}a_{ijk}b_{ijk}$$? I'm not familiar with tensor product so that also contributes to my struggle here.

I'm learning this to solve this problem of mine.

Jean-Claude Arbaut
  • 23,601
  • 7
  • 53
  • 88
LWZ
  • 295

1 Answers1

6

The result is a new array c, with $$ c_{kl} = \sum_{i,j} a_{ijk} b_{jil} . $$

Hans Lundmark
  • 55,038
  • Thanks! And I found np.einsum('ijk,jil', a, b) also gives the same result, without the ->kl part. – LWZ Mar 21 '13 at 15:10
  • 2
    The reason for this is that what comes after -> is how the output is treated. As you sum along repeated indices $i$ and $j$, the output depends on $k$ and $l$, the default without the -> is to keep it this way. If you want to sum over those indices as well (sum the final columns or rows), you can do it like this np.einsum('ijk,jil->k', a, b). If you want to sum columns and rows, you could do np.einsum('ijk,jil->', a, b). – Ramon Crehuet Jan 11 '16 at 20:52
  • numpy.einsum hits superfast BLAS-like routines? – a06e Nov 29 '19 at 14:56