Numpy库中einsum函数用法

【导读】 einsum全称Einstein summation convention（爱因斯坦求和约定），又称为爱因斯坦标记法。能够计算任何维度的张量收缩。einsum的写法省去了求和符号，显得更加简洁。

一、一维张量收缩

对于一维张量，也即向量（Vector）。其收缩为零维张量，也即标量（Scalar）：

$$c=\sum _i{a}_i{b}_i$$所以einsum 的写法就是：
$$c=a_i{b}_i$$用代码表示为：

>>> a = np.arange(10)
>>> b = np.arange(10)+1
>>> a
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> b
array([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10])
>>> np.einsum('i,i', a, b)
330

二、二维张量收缩

同样对于二维张量，也即矩阵（Matrix）。其收缩为一维张量或零维张量：

2.1 收缩到零维张量

$$c=a_{ij}{a}_{ij}=\sum _{i,j}{a}_{ij}{a}_{ij}$$也即求矩阵内积

array([[0, 1],
       [2, 3],
       [4, 5]])
>>> np.einsum('ij,ij', a, b)
70

2.2 收缩到一维张量

$$c_i=\sum _j{A}_{ij}{B}_{ij}$$或$$c_j=\sum _i{A}_{ij}{B}_{ij}$$
其中，$c_i$为两矩阵对应位置元素相乘后再把列元素相加；$c_j$为两矩阵对应位置元素相乘后把行相加

>>> np.einsum('ij,ij->i', a, b)
array([ 8, 62])
>>> np.einsum('ij,ij->j', a, b)
array([12, 22, 36])

再给出一个例子

>>> a1 = np.array([[0,1],[2,3],[4,5]])
>>> a1
array([[0, 1],
       [2, 3],
       [4, 5]])
>>> np.einsum('ji,ij->i', a1, a)
array([10, 40])

再给出一个矩阵与向量的例子

>>> a2 = np.array([[1,2],[3,4],[5,6]])
>>> b2 = np.array([1,2])
>
>>> np.einsum('ij,j', a2, b2)
array([ 5, 11, 17])
>>> np.einsum('ij,j->...', a2, b2)	# 收缩到零维
33

再给出一个多个向量的例子

>>> c2 = np.array([1,2,3])
>>> np.einsum('ij,j,i->i', a2, b2, c2)
array([ 5, 22, 51])
>>> np.einsum('ij,j,i->...', a2, b2, c2)  # 收缩到零维
78

三、三维张量收缩（重难点）

由于其组合非常多样，如三维张量与矩阵、向量的爱因斯坦求和运算，收缩到二维、一维、零维，我们不一一讨论。这里我们只给出三维张量之间进行爱因斯坦求和、收缩到二维张量的例子

3.1 例1

>>> a = np.arange(6).reshape(1,2,3)
>>> b = np.arange(24).reshape(2,3,4)
>>> a
array([[[0, 1, 2],
        [3, 4, 5]]])
>>> b
array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]],

       [[12, 13, 14, 15],
        [16, 17, 18, 19],
        [20, 21, 22, 23]]])
        
>>> np.einsum('ijk,jkl->il', a, b)
array([[220