Numpy ndarray 乘法

Question

我有两个3D的 numpy ndarray

A=np.array([[[1, 1],
             [1, 1],
             [1, 1]],

            [[2, 2],
             [2, 2],
             [2, 2]]])

B=np.array([[[ 2,  0],
             [ 0,  2]],

            [[ 2, -2],
             [-2,  2]]])

我想创建一个AB数组，其中元素AB_ijk=∑_m (A_ijm*B_imk)，其中求和仅针对m-index（重复），而不针对i（反过来重复）。

换句话说，我可以使用此for循环获得di AB ndarray

for i in range(2):
    AB[i,:,:]=np.dot(A[i,:,:],B[i,:,:])

AB等于

array([[[ 2.,  2.],
    [ 2.,  2.],
    [ 2.,  2.]],

   [[ 0.,  0.],
    [ 0.,  0.],
    [ 0.,  0.]]])

有没有避免使用for循环的方法？我如何使用tensordot或者einsum获取AB数组？

感谢您的回答，非常感激。

- musine

Answer 1

在足够新的NumPy（1.10+）上，您可以执行以下操作

AB = np.matmul(A, B)

或（如果您还安装了Python 3.5+）：

AB = A @ B

如果您没有NumPy 1.10+，您可以这样做：

AB = np.einsum('ijm,imk->ijk', A, B)

对于较大的J/M/K维度，特别是如果您拥有良好的BLAS，考虑使用显式的for循环和dot也许更值得。BLAS矩阵乘法可能比更多解释性Python代码的开销节省更多时间。我认为np.matmul和@应该利用dot所做的相同事情，但我不认为np.einsum会这样做。

- user2357112

---

确认：einsum 目前尚未使用 BLAS。 - Daniel

Answer 2

ABijk=∑m (Aijm*Bimk) 的翻译是：

AB = np.einsum('ijm,imk->ijk', A, B)

我认为matmul运算符也可以处理这个问题。

AB = A @ B

因为它将正常的点应用到最后的2个维度上，其余维度则被携带为免费行李。

测试一下并告诉我是否可用。

- hpaulj