numpy中张量的矩阵向量乘法

Question

我有一个形状为(l,l) 的numpy.array A 和另一个形状为(l,m,n) 的numpy.array B。通常，B 中的第二个和第三个维度对应于空间单元格，而第一个维度对应于其他内容。

我要计算

l,m,n = 2,3,4                 # dummy dimensions
A = np.random.rand(l,l)       # dummy data
B = np.random.rand(l,m,n)     # dummy data
C = np.zeros((l,m,n))


for i in range(m):
  for j in range(n):
    C[:,i,j] = A@B[:,i,j]

即，在每个空间单元中，我想执行矩阵向量乘法。 由于我必须经常这样做，我想知道是否有更紧凑的方法可以用numpy 编写这个。 （特别是因为有几种情况张量的形状为(l,m,n,o,p)。）

提前谢谢你！

Answer 1

我使用np.einsum找到了答案：

np.einsum('ij,jkl->ikl', A,B)

说明： 爱因斯坦符号意味着我们对匹配的下标求和。

np.einsum('ij,jkl->ikl', A,B)
= rewritten in math terms
A_{i,j} B_{j,k,l}
= Einstein notation implies summation
sum_j A_{i,j} B_{j,k,l}