如何在numpy下进行高维矩阵乘法？

Question

在 numpy 下我想执行“通常的”矩阵乘法，如下所示：

C=A*B

哪里

A 是一个“2D 类”矩阵，但每个矩阵元素的形状为 (1,5)

和

B 是一个“1D-kind”向量，但每个向量元素具有形状 (20,5)

结果

C 应是“一维类”向量，但每个向量元素再次具有形状 (20,5)

我尝试手动生成 C 的元素 C1 和 C2：

>>> A.shape
(2, 2, 1, 5)
>>> B.shape
(2, 20, 5)
>>> C0 = A[0,0]*B[0]+A[0,1]*B[1]
>>> C0.shape
(20, 5)
>>> C1 = A[1,0]*B[0]+A[1,1]*B[1]
>>> C1.shape
(20, 5)
>>>

将 A 的 (1,5) 广播到 B 的 (20,5) 可以按预期工作。

但是，我无法找出如何将其写成矩阵乘法：

C = np.matmul(A, B)

当然这是行不通的，因为 numpy 不知道我想要对哪些索引求和。但我想一定存在一些简单的“numpythonic”解决方案......

Answer 1

您可以使用

einsum

：

C = np.einsum('ij...,j...->i...', A, B)

这是 A 的第二个维度与 B 的第一个维度的收缩。其余维度将继续运行。