Python中向量与张量的矩阵乘法：寻求单行实现方案

你可以利用NumPy的广播机制、批量矩阵乘法或爱因斯坦求和（einsum）来实现单行代码替代循环，以下是几种可行方案：

方案1：批量矩阵乘法（matmul） #

利用np.matmul对高维张量的批量处理能力，先将向量M扩展为二维行向量，再与张量d做批量矩阵乘法，最后压缩冗余维度：

a = G * np.matmul(M[np.newaxis, :], d).squeeze(axis=1)

• 原理：M[np.newaxis, :]将原N维向量转为形状(1, N)的行向量；np.matmul会自动广播行向量到与d的第一个维度匹配（即n次重复），对每个d[i,:,:]执行矩阵乘法，得到形状(n, 1, D)的结果；squeeze(axis=1)去掉中间的冗余维度，最终得到(n, D)的矩阵。

方案2：爱因斯坦求和（einsum） #

用np.einsum直接描述运算逻辑，可读性更强：

a = G * np.einsum('m, imd->id', M, d)

• 原理：下标'm, imd->id'表示：对每个i（对应d的第一个维度）和每个d（对应d的第三个维度），求和M[m]与d[i,m,d]的乘积，完全等价于循环中M @ d[i,:,:]的运算，直接输出(n, D)的结果。

方案3：广播+求和 #

通过广播实现元素级相乘，再对矩阵乘法的求和维度进行汇总：

a = G * (M[np.newaxis, :, np.newaxis] * d).sum(axis=1)

• 原理：M[np.newaxis, :, np.newaxis]将向量转为形状(1, N, 1)的张量，与d(n, N, D)广播相乘后得到(n, N, D)的元素乘积张量；对axis=1（即N维度）求和，就得到行向量乘矩阵的结果，最终形状为(n, D)。

注：如果你的张量d实际维度是N×N×D（而非描述中的n×N×D），只需将上述代码中的维度下标对应调整即可，核心逻辑不变。

内容的提问来源于stack exchange，提问作者william paine