如何将(M,N,N)的3D numpy数组与(M,N)的2D数组相乘得到(M,N)数组？

你需要对每个批次的矩阵-向量对执行乘法，以下几种方法都可以实现需求：

方法1：利用np.matmul（或@运算符）的广播特性 #

np.matmul支持广播机制，只需给数组B添加一个末尾维度，让其形状变为(M,N,1)，就能和形状为(M,N,N)的A进行广播式矩阵乘法，最后去掉多余维度即可：

import numpy as np

M, N = 4, 3
A = np.ones((M,N,N))
B = np.array([[1,1,1],[1,0,0],[0,0,0],[1,2,3]])

# 添加维度后执行矩阵乘法，再压缩最后一维
result = np.matmul(A, B[..., np.newaxis]).squeeze(-1)
# 或者用@运算符，效果完全一致
# result = (A @ B[..., np.newaxis]).squeeze(-1)

print(result)

输出：

[[3. 3. 3.]
 [1. 1. 1.]
 [0. 0. 0.]
 [6. 6. 6.]]

方法2：使用np.einsum明确指定维度运算 #

你之前尝试einsum失败应该是维度下标书写错误，正确的下标表达式是"mij,mj->mi"，各下标对应：

• m：批次维度（M）
• i：矩阵的行维度（第一个N）
• j：矩阵的列/向量的维度（第二个N）
  通过收缩j维度，得到每个批次下的N维结果：

result = np.einsum('mij,mj->mi', A, B)
print(result)

输出和方法1完全一致。

方法3：利用np.dot结合形状变换 #

把A展开成(M*N, N)的二维数组，B转置为(N,M)，执行点积后再恢复原批次形状：

result = np.dot(A.reshape(-1, N), B.T).reshape(M, N)
print(result)

同样能得到期望的结果。

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