【学习笔记】einsum函数

einsum全称Einstein summation convention（爱因斯坦求和约定），又称为爱因斯坦标记法。einsum的写法省去了求和符号，显得更加简洁。

常规	einsum
$c={\sum }_i{a}_i{b}_i$	$c=a_i{b}_i$
$c_j={\sum }_j{A}_{ij}{b}_j$	$c_j=A_{ij}{b}_j$
$c_{klmn}={\sum }_i{\sum }_j{a}_{ijkl}{b}_{ijmn}$	$a_{ijkl}{b}_{ijmn}$

函数使用

多维张量求和

实现$c_{jk}={\sum }_i{a}_{ijk}$

c = a.sum(axis=0)           # 常规
c = np.einsum('ijk->jk', a) # einsum

更高维向量的求和，实现$c_{\ast }={\sum }_i{a}_{i\ast }$，可以直接将剩下的维度替换成省略号。

c = np.einsum('i...->...', a) # einsum

矩阵乘法

实现$c_{ik}={\sum }_j{a}_{ij}{b}_{jk}$

c = np.dot(a, b)                 # 常规
c = np.einsum('ij,jk->ik', a, b) # einsum

更高维张量的乘法，实现$c_{kl}={\sum }_i{\sum }_j{a}_{ijk}{b}_{ijl}$

c = np.einsum('ijk,jkl->kl', a, b) # einsum