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

Question

在 numpy 中有一个操作我发现如果不循环就很难实现

那个操作是我有两个输入：beta,x

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

我感兴趣的操作可以使用for循环完成如下

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_

在这里让我可以在没有循环的情况下使用技巧，以便整个操作可以高效计算？

Answer 1

您有兴趣在公共 K 维度上乘以元素，并将结果保持在其余维度上。

这意味着您可以使用 np.einsum 使用 beta、x 的尺寸以及您感兴趣的形状结果，例如 '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

不是问题的一部分，但在谈论 np.einsum...

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