如何向量化实现NumPy多维数组指定轴的矩阵-向量乘法运算

方法1：使用np.einsum（最简洁，无需调整维度） #

你之前用einsum没得到正确结果是索引设置错误，正确写法如下：

b = np.einsum('mnij,mnj->mni', A, x)

索引规则说明：

• mnij对应4维数组A的四个轴：批次M、批次N、矩阵行D、矩阵列D
• mnj对应3维数组x的三个轴：批次M、批次N、向量维度D
• 输出mni保留两个批次轴，结果维度为(M,N,D)，完全符合你的需求

方法2：使用广播版np.matmul #

利用numpy matmul对批次维度的自动广播特性，只需给x增加一个列向量维度，运算后删除多余维度即可：

# x[..., np.newaxis]把x形状从(M,N,D)转为(M,N,D,1)，squeeze删除最后一维的1
b = np.matmul(A, x[..., np.newaxis]).squeeze(axis=-1)

关于tensordot形状不匹配的说明 #

tensordot默认不会自动对齐前两个批次维度，除了你指定的求和轴外，其余轴会做外积，因此会出现形状不匹配的问题，该场景下不推荐使用tensordot实现，上述两种方法效率和可读性都更优。

正确性验证测试 #

import numpy as np
# 测试用例：预期输出为[5,11]
A = np.array([[[[1,2],[3,4]]]]) # 形状(1,1,2,2)
x = np.array([[[1,2]]]) # 形状(1,1,2)

print(np.einsum('mnij,mnj->mni', A, x))
# 输出：[[[ 5 11]]]，符合预期
print(np.matmul(A, x[..., None]).squeeze(-1))
# 输出：[[[ 5 11]]]，符合预期

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