python - Numpy 3d 数组矩阵乘法函数

Question

假设我有一个 ndarray，形状为 (m,n,n) 的 W 和一个维度为 (m,n) 的向量 C。我需要按以下方式将这两个相乘

result = np.empty(m,n)
for i in range(m):
    result[i] = W[i] @ C[i]

我如何在没有循环的情况下以矢量化的方式做到这一点？

Answer 1

因为，您需要保持 W 和 C 的第一个轴对齐，同时通过矩阵乘法失去它们的最后一个轴，我建议使用np.einsum对于一个非常有效的方法，就像这样 -

np.einsum('ijk,ik->ij',W,C)

np.tensordot 或 np.dot 没有保持轴对齐的功能，而这正是 np.einsum 改进的地方。