NumPy中einsum使用笔记

Python的Numpy提供了很多高效的科学计算函数，einsum便是其中一个。einsum函数非常灵活，也非常高效，运行时占用的存储空间很小。由于其强大的表现力和智能循环，它在速度和内存效率方面通常可以超越我们常见的array函数。但缺点是，可能需要一段时间才能理解符号，有时需要尝试才能将其正确的应用于棘手的问题。

1. 爱因斯坦求和约定

在线性代数中，我们最多涉及的是二阶及以下的张量．在这种情况下，纸面上可以很方便地写出低阶张量的矩阵形式，高阶的张量，它们的坐标就没法用矩阵表示．我们当然可以把矩阵拓展为立体阵等概念，但随着阶数上升，这种表示法的复杂程度几何级增加；我们也可以使用张量词条中所提过的向量矩阵的方法，比起立体阵要清楚一些，但套娃式的表达方式也对理解一个张量的性质造成了障碍．

爱因斯坦求和约定正是为了简洁地表达高阶张量的坐标运算而存在的．

假设 矩阵大小分别是 2*3 和 3*2，矩阵乘法的定义如下：

爱因斯坦求和是一种对求和公式简洁高效的记法，其原则是当变量下标重复出现时，即可省略繁琐的求和符号。

比如求和公式：

其中变量 a 和变量 b 的下标重复出现，则可将其表示为：

同理

进一步地，我们可以得到矩阵乘法的一个抽象 

 

2. einsum的用途

einsum方法正是利用了爱因斯坦求和简洁高效的表示方法，从而可以驾驭任何复杂的矩阵计算操作。基本的框架如下：

以计算下式为例

可以写成如下形式

C = einsum('ij,jk->ik', A, B)

上述操作表示矩阵A与矩阵B的点积。输入的参数分为两部分:

• 前面表示计算操作的指令串，
• 后面是以逗号隔开的操作对象（数量需与前面对应）。

其中在计算操作表示中，

• "->" 左边是以逗号隔开的下标索引，重复出现的索引即是需要爱因斯坦求和的；
• "->" 右边是最后输出的结果形式。

在矩阵之间的运算中，下标可以分为两类：

• 自由标(Free index)，也就是在输入和输出端都出现的下标
• 哑标(Summation index)，在输入端出现但输出端没有出现的下标

实际上，即使是这个简单的例子，我测试使用 einsum 计算的速度要是常规手段的 3 倍。

用图像表示上述过程

A = np.array([[1, 1, 1],
              [2, 2, 2],
              [5, 5, 5]])
 
B = np.array([[0, 1, 0],
              [1, 1, 0],
              [1, 1, 1]])

我们将轴标签的输入/输出轴画出来：

要理解上述计算过程，记住下面这三条规则：

• 输入arrays沿着重复的那个轴做乘法计算。

在本例中，轴标签的输入参数ij,jk中 'j' 被重复用了 2 次：1次表示A的第二个轴，1次表示B的第一个轴。这意味着要计算A的第二个轴（行方向）与B的第一个轴（列方向）的乘积。所以，此时要保证输入合法，就需要保证A的行长与B的列长一致。

• 如果某个轴标签在输出标签中消失了，则表示要沿着该轴做求和计算。

此处，'j' 没有出现在输出标签中，这表示在执行乘法运算后，需要再沿着'j'轴做求和运算，求和运算减少了输出 array 的 1 个纬度。如果我们将输出标签改为'ijk'，那么因为少了求和运算，最终得到的输出 array 将会是 3x3x3 形状的。

如果此处我们将输入/输出标签改为'ij,jk->'，也即令所有的标签都不出现在输出轴标签里，那么我们将得到一个标量，这个标量是所有元素的和。

• 我们可以获取任意顺序轴的结果。

如果我们在轴标签中不写'->'，那么 numpy 会将只出现一次的标签按照字母顺序组合，作为输出轴标签，所以 ij,jk->ik 和 ij,jk 效果上是等价的。指定轴顺序的输出，可以通过指定轴标签的顺序获得。例如，'ij,jk->ki'可以得到'ij,jk->ik'的转置矩阵。

了解上面三条基本规则后，再来看einsum如何计算矩阵乘法的就简单了。下图是左边是计算np.einsum('ij,jk->ijk', A, B)的结果，右图则是按照'j'轴求和后的结果：

按照前文所述，einsum是非常节约空间的计算函数，所以对于np.einsum('ij,jk->ik', A, B)，einsum并不构建临时的 3D array 然后求和，而是直接在 2D 空间累加得到最终结果。

3. einsum的简单操作

下面两张表描述了 einsum 的基本操作。

若 A 和 B 为两个 1D arrays（假设在相应的操作中，A和B的形状总是合适的），那么：

若 A 和 B 为两个 2D arrays（假设在相应的操作中，A和B的形状总是合适的），那么：

4. einsum轴标签中的'...'符号

在处理比较多的纬度时，为了方便，可以像 numpy array 一样使用 '...' 符号省略一些纬度的显式表示。例如，

np.einsum('...ij,ji->...', a, b)

这一行代码计算的是 a 的后两个轴与 2D array b 的乘积。

5. 注意事项

einsum 在求和时，不会提升数据类型。如果我们使用了位宽比较小的数据类型，可能会得到不期望的结果：

>>> a = np.ones(300, dtype=np.int8)
>>> np.sum(a) # correct result
300
>>> np.einsum('i->', a) # produces incorrect result
44

另外，einsum在 numpy 的计算库中并不一定总是最快的。类似于 dot和 inner 的函数一般链接到快速计算库 BLAS，可能会快于 einsum。

参考文献

einsum方法详解（爱因斯坦求和） - 知乎

nPython计算库NumPy的einsum（爱因斯坦和）使用简介 - 刘冲的博