将 numpy.einsum 用于切片的张量点

Question

I have an array of shape N, j, j and another of shape M, j, j . I want to calculate their tensor dot to eventually have a matrix where the i,j entry is np.tensordot(arr1[i, :, :], arr2[j, :, :]) I've tried looping but it is ridiculously slow, I've read about np.einsum but unfortunately cannot figure it out no matter how much I read. My most recent attempt; np.einsum('ilk,ium->lu', arr1, arr2) But I keep getting errors that the shapes can't be broadcasted. Would appreciate any pointers, thanks!

example code:

  arr1 = np.zeros((5, 2, 2))
  arr2 = np.zeros((4, 2, 2))
  arr2[1,:,:][1,1] = 2
  arr1[1,:,:][1,1] = 3
  np.tensordot(arr1[1,:,:], arr2[1,:,:])

in this case, the tensor dot would give me 6. That is what I am interested in, for each i,j .

Answer 1

In [41]: x=np.arange(2*3*3).reshape(2,3,3)
In [42]: y=np.arange(4*3*3).reshape(4,3,3)

double contraction on the last 2 dim:

In [43]: np.einsum('ikl,jkl->ij',x,y)
Out[43]: 
array([[ 204,  528,  852, 1176],
       [ 528, 1581, 2634, 3687]])

test one value:

In [44]: np.tensordot(x[0],y[0])
Out[44]: array(204)

Same thing dot (and extra dimension)

In [47]: np.dot(x.reshape(-1,9),y.reshape(-1,9,1))
Out[47]: 
array([[[ 204],
        [ 528],
        [ 852],
        [1176]],

       [[ 528],
        [1581],
        [2634],
        [3687]]])

np.tensordot with the various axis options can be a bit tricky to use. One way or other it reshapes and transposes the arrays so it can call np.dot . Then it may do some further manipulation.

Or using broadcasting and multiaxis sum:

In [48]: (x[:,None,:,:]*y[None,:,:,:]).sum(axis=(2,3))
Out[48]: 
array([[ 204,  528,  852, 1176],
       [ 528, 1581, 2634, 3687]])

and a non-loop tensordot:

In [50]: np.tensordot(x,y,axes=((1,2),(1,2)))
Out[50]: 
array([[ 204,  528,  852, 1176],
       [ 528, 1581, 2634, 3687]])