沿多维数组指定轴的矢量化矩阵-矢量乘法

Question

假设我有一个大小为 (M,N,D,D) 的数组 A 和一个大小为 (M,N,D) 的向量 x。我想生成大小为 (M,N,D) 的数组 b，其中 b[m,n,:] = np.matmul(A[m,n,:,:],x[m,n,: ]).所以基本上，A[m,n,:] 中包含的 D×D 矩阵和 x[m,n,:] 中包含的大小为 D 的向量应该按照标准矩阵向量乘法相乘。一个天真的 for 循环实现是

for m in range(M):
    for n in range(N):
        b[m,n,:] = np.matmul(A[m,n,:],x[m,n,:])

显然这对于大的 N 和 M 来说是低效的（我就是这种情况；N 的数量级为 100,000，M 的数量级为 1000）。我的问题是：如何向量化这个循环？我试过使用 np.tensordot 但无论我如何尝试构造它，我总是得到形状不匹配，大概是因为 A 有 4 个轴而 x 有 3 个（请注意，据我所知，np.tensordot 是正确使用的功能，我只是没有正确使用它）

编辑尝试使用 np.einsum，但我没有得到预期的结果（应该是 [5,11]：

Answer 1

这就是 np.einsum 的亮点：

np.einsum('mnij,mnj->mni', A,B)

检查

np.allclose(np.einsum('mnij,mnj->mni', A,x), b)

returns True

Answer 2

广播也可以达到同样的效果：

x1 = x.reshape(M,N,d,1)
b1 = np.matmul(A,x1)
b2 = b1.reshape(M,N,d)

沿多维数组指定轴的矢量化矩阵-矢量乘法

Vectorized Matrix-Vector multiplication along specified axis of multidimensional array

arrays

numpy

vectorization

scipy

multidimensional-array