矩阵向量乘法与 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}