python - Python：在 numpy 数组上进行矢量化

Question

X并且Y都是尺寸为 的 3d 数组(a,b,c)。我的目标是做一个点积。

考虑索引i和j是标量的情况，并且(X[i,:,j].T).dot(Y[i,:,j])很简单并返回一个标量。

但是，如果我尝试进行矢量化，i并j成为一维数组，并(X[i,:,j].T).dot(Y[i,:,j])返回一个矩阵，但我期待一个一维数组作为结果。我该如何解决这个问题？

标签： pythonarraysnumpy

Answer 1

使用列表理解的幼稚实现：

a,b,c = X.shape

r1 = [(X[i,:,j].T).dot(Y[i,:,j]) for i in range(a) for j in range(c)]

使用np.einsum实现：

r2 = np.einsum('ijk,ijk->ik', X,Y).flatten()