operation of Einstein sum of 3D matrices

Question

The following code indicates that the Einstein sum of two 3D (2x2x2) matrices is a 4D (2x2x2x2) matrix.

$ c_{ijlm} = \Sigma_k a_{i,j,k}b_{k,l,m} $

$ c_{0,0,0,0} = \Sigma_k a_{0,0,k}b_{k,0,0} = 1x9 + 5x11 = 64 $

But, c_{0,0,0,0} = 35 according to the result below:

>>> a=np.array([[[1,2],[3,4]],[[5,6],[7,8]]])
>>> b=np.array([[[9,10],[11,12]],[[13,14],[15,16]]])
>>> c=np.einsum('ijk,klm->ijlm', a,b)
>>> c
array([[[[ 35,  38],
         [ 41,  44]],

        [[ 79,  86],
         [ 93, 100]]],


       [[[123, 134],
         [145, 156]],

        [[167, 182],
         [197, 212]]]])

Could someone explain how the operation is carried out?

score 1 · Accepted Answer · answered Jan 16 '21 at 18:21

The particular element that you are testing, [0,0,0,0] is calculated with:

In [167]: a[0,0,:]*b[:,0,0]
Out[167]: array([ 9, 26])
In [168]: a[0,0,:]
Out[168]: array([1, 2])
In [169]: b[:,0,0]
Out[169]: array([ 9, 13])

It may be easier to understand if we reshape both arrays to 2d:

In [170]: A=a.reshape(-1,2); B=b.reshape(2,-1)
In [171]: A
Out[171]: 
array([[1, 2],
       [3, 4],
       [5, 6],
       [7, 8]])
In [172]: B
Out[172]: 
array([[ 9, 10, 11, 12],
       [13, 14, 15, 16]])
In [173]: A@B
Out[173]: 
array([[ 35,  38,  41,  44],
       [ 79,  86,  93, 100],
       [123, 134, 145, 156],
       [167, 182, 197, 212]])

The same numbers, but in (4,4) instead of (2,2,2,2). It's easier to read the (1,2) and (9,13) off of A and B.

Thank you for the illustration. forgot that the depth is the first dimension, rather than the third. btw, I think use c0000=np.dot(a[0,0,:], b[:,0,0]) would be simpler. — techie11, Jan 16 '21 at 23:20

operation of Einstein sum of 3D matrices

1 Answers1