如何在不使用 for 循环的情况下对这两个 numpy 操作进行矢量化？

Question

there is an operation in numpy I've found hard to implement without looping

That operation is that I have two inputs: beta,x

beta.shape = (M,N,K) and x.shape = (I,K)

The operation I'm interested in can be done using for loop as follows

result = np.zeros((M,N,I,K)) # buffer to save my operation results
for m in range(M):
    for n in range(N):
         beta_ = beta[m][n] # has shape (K,)
         result[m][n] = x * beta_

Let tricks here that I can do without loops so that the whole operation can be computational efficient?

Answer 1

You're interested in multiplying the elements across the common K dimension, and keeping the results along the remaining dimensions.

This means you can use np.einsum using the dimensions of beta, x, and the shape result you're interested in like 'mnk,ik->mnik' :

import numpy as np

M = 4
N = 3
I = 7
K = 6

beta = np.arange(M*N*K).reshape(M,N,K)
x = np.arange(I*K).reshape(I,K)

result1 = np.zeros((M,N,I,K)) # buffer to save my operation results
for m in range(M):
    for n in range(N):
         beta_ = beta[m][n] # has shape (K,)
         result1[m][n] = x * beta_


result2 = np.einsum('mnk,ik->mnik', beta, x)

print (np.array_equal(result1,result2))

 True

Not part of the question but while talking about np.einsum ... If you want to sum across any of these dimensions, you can omit it from the end dimensions:

import numpy as np

M = 4
N = 3
I = 7
K = 6

beta = np.arange(M*N*K).reshape(M,N,K)
x = np.arange(I*K).reshape(I,K)

result1 = np.zeros((M,N,I,K)) # buffer to save my operation results
for m in range(M):
    for n in range(N):
         beta_ = beta[m][n] # has shape (K,)
         result1[m][n] = x * beta_
result1 = result1.sum(axis=1)


result2 = np.einsum('mnk,ik->mik', beta, x)

print (np.array_equal(result1,result2))

 True