如何对多矩阵乘法做向量化处理，numpy.einsum的正确语法是什么

结论 #

你给出的einsum写法np.einsum('ij,jki->ik', A, B)是完全正确的，和你提供的循环实现逻辑完全等价。

规则对应解释 #

我们可以对应你的运算场景拆解einsum字符串的含义：

• ij 对应二维矩阵A的两个维度：第一个维度i是样本序号（对应循环里的遍历变量，范围0_{999），第二个维度`j`是特征维度（范围0}89）
• jki 对应三维矩阵B的三个维度：第一个维度j和A的特征维度对齐，做矩阵乘法的内积运算；第二个维度k是输出结果的特征维度；第三个维度i和A的样本序号对应，保证每个样本i都取B中对应i位置的二维切片参与计算
• ->ik 表示最终输出保留样本维度i和结果特征维度k，输出矩阵shape刚好是你需要的1000*90，和目标矩阵C的维度完全匹配。

验证与优化建议 #

• 你可以用小尺寸随机矩阵验证结果一致性：生成小维度的测试矩阵（比如A = np.random.rand(3,2)，B = np.random.rand(2,2,3)），分别用循环和上述einsum写法计算结果，再调用np.allclose(循环计算结果, einsum计算结果)，返回True即可确认正确性。
• 如果要进一步提升大尺寸矩阵的计算效率，可以给einsum添加optimize='optimal'参数，写法为np.einsum('ij,jki->ik', A, B, optimize='optimal')，numpy会自动选择最优的计算路径，降低运算开销。

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