如何使用np.einsum编写多索引矩阵相乘的求和表达式？

两矩阵求和公式实现 #

对应公式 f(i) = \sum_{j,k} a_ij * b_jk，最终仅保留i维度，对j、k两个维度求和，einsum写法如下：

np.einsum('ij,jk->i', A, B)

参数说明：

• ij 匹配m×n矩阵A的两个维度，i对应行维度、j对应列维度
• jk 匹配n×r矩阵B的两个维度，j对应行维度、k对应列维度
• 箭头后仅保留i，所有未出现在箭头后的下标j、k会自动执行求和操作，最终输出形状为(m,)，和公式定义完全一致。

新增矩阵C后的写法调整 #

更新后的公式新增r×p的矩阵C参与计算，需要对j、k、l三个维度求和，仅保留i维度，仅需要做两处调整：

1. 在下标规则字符串中新增矩阵C对应的下标段kl（k对应C的行维度、l对应C的列维度）
2. 在入参列表中新增矩阵C
   最终实现代码如下：

np.einsum('ij,jk,kl->i', A, B, C)

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