Python Numpy 沿第二维和第三维向量化

Question

I want to vectorize the calculation along the second and third axis, but I cannot find a way to deal with the last for loop:

n = 3000
D = 2000
delta = np.arrange(D * n).reshape(D, n)
M = np.arrange(D * n * n).reshape(D, n, n)
result = np.zeros(D, n)
for i in range(D):
    result[i, :] = np.dot(delta[i, :], M[i, :, :])

Answer 1

Let's construct a small example (no need to use big sizes like 3000):

In [122]: n,D = 3,2                                                             
In [123]: delta = np.arange(D*n).reshape(D,n)    # not arrange                                   
In [124]: M = np.arange(D*n*n).reshape(D,n,n)                                   
In [125]: result = np.zeros(D, n) 
     ...: for i in range(D): 
     ...:     result[i, :] = np.dot(delta[i, :], M[i, :, :]) 
     ...:                                                                       
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-125-413de80d25a5> in <module>
----> 1 result = np.zeros(D, n)
      2 for i in range(D):
      3     result[i, :] = np.dot(delta[i, :], M[i, :, :])
      4 

TypeError: data type not understood

Oops, your np.zeros call is wrong!

In [126]: result = np.zeros((D, n),int) 
     ...: for i in range(D): 
     ...:     result[i, :] = np.dot(delta[i, :], M[i, :, :]) 
     ...:                                                                       
In [127]: result                                                                
Out[127]: 
array([[ 15,  18,  21],
       [150, 162, 174]])

einsum is a handy way of matching diverse dimensions. Note how the i dimension matches yours.

In [128]: np.einsum('ij,ijk->ik',delta, M)                                      
Out[128]: 
array([[ 15,  18,  21],
       [150, 162, 174]])

Let's try the matmul/@ function which is supposed to handle 'batch' dots like yours:

In [129]: np.matmul(delta, M)                                                   
Out[129]: 
array([[[ 15,  18,  21],
        [ 42,  54,  66]],

       [[ 42,  45,  48],
        [150, 162, 174]]])
In [130]: _.shape                                                               
Out[130]: (2, 2, 3)

Oops, too much. Let's add a dimension to delta to make it (D,1,n). That way the D dimension is consistently the first of 3:

In [131]: np.matmul(delta[:,None,:], M)                                         
Out[131]: 
array([[[ 15,  18,  21]],

       [[150, 162, 174]]])

Those are the correct values, but the shape is (2,1,3). We can get rid of that middle dimension with reshape or index or squeeze:

In [132]: np.matmul(delta[:,None,:], M).squeeze()                               
Out[132]: 
array([[ 15,  18,  21],
       [150, 162, 174]])