1 背景
pyTorch和numpy等库都提供了einsum函数函数。einsum，全称Einstein summation convention（爱因斯坦求和约定），是一个方便而强大的工具。
我们熟悉的矩阵乘法、转置、求矩阵的迹等操作都可视为einsum的特例。除此之外，它还可以表示更复杂的运算，例如将矩阵乘法、转置等操作复合在一起表示为一个操作。
先从矩阵乘法说起。我们可以这样用einsum表示矩阵乘法，
import torch
a = torch.rand(2, 2)
b = torch.rand(2, 2)
c = torch.einsum('ij,jk->ik', a, b)在ij,jk->ik这个式子中，j在->左侧重复使用，且没有出现在->右侧，因此我们需要对j求和。由此可知einsum('ij,jk->ik', a, b)的实际含义为\(c_{ik}=\sum_j a_{ij} b_{jk}\)，等同于矩阵乘法。
下面我们简单验证一下计算的正确性。
eps = 1e-5
assert (c - a @ b).abs().mean() < eps
print('检查通过')检查通过上面的代码中，a @ b是pytorch中计算a和b矩阵乘法时的一般写法。可以看到我们用einsum得到的计算结果与标准写法的计算结果是一致的。
代码torch.einsum('ij,jk->ik', a, b)使用了->符号分隔输入和输出。这被称为显式的einsum。你也可以不使用->，只给出输入的索引，这被称为隐式的einsum。例如torch.einsum('ij,jk', a, b). 注意因为隐式的einsum不指定输出的索引，因此输出的索引将会按字母表顺序排列。
为了简便起见。这篇文章里我只讨论显式的einsum。基于以上介绍的基础规则，我想在这篇文章中分享一些有趣的发现和思考：
- 顺序无关性：einsum的输入参数的顺序是可交换的；
- 手写einsum函数：如何基于python实现一个基础的einsum函数；
- einsum的梯度：einsum函数的梯度仍然可以用einsum表示；
2 顺序无关性
einsum函数的要点是：1. 根据索引从输入参数中取数，计算累积的乘积；2. 对重名的，在输出中没有出现的索引求和。
可以看到einsum的规则和输入参数的顺序没有关系。因此，einsum函数的输入参数是可以交换顺序的。
例如，虽然矩阵乘法没有交换律，但是基于einsum计算矩阵乘法时，你可以交换参数的顺序。
for _ in range(100):
a = torch.rand(2, 2)
b = torch.rand(2, 2)
c = torch.einsum('ij,jk->ik', a, b) # 正序
d = torch.einsum('jk,ij->ik', b, a) # 反序
assert (c - d).abs().mean() < eps
print('检查通过')检查通过3 手写einsum
einsum函数的计算是这样一个过程：一个遍历所有索引，计算输入元素的乘积，然后求和的过程。如果我们有一个函数能够遍历所有的索引，那会很方便。
我们先来实现这样的一个函数iter_elements：
import torch
def iter_elements(sizes):
if len(sizes) == 0:
yield []
return
for i in range(sizes[0]):
for j in iter_elements(sizes[1:]):
yield (i, *j)
for idx in iter_elements([2, 3]):
print(idx)(0, 0)
(0, 1)
(0, 2)
(1, 0)
(1, 1)
(1, 2)上面的代码实现了一个遍历所有可能索引的函数iter_elements. 对于尺寸为\(2\times 3\)的矩阵，这个函数输出了所有可能的六个索引。这是一个非常方便的函数。接下来我们就基于它在einsum函数中遍历每一个元素。
def einsum_forward(equation, *operands):
input_dims, output_dim = equation.split('->')
input_dims = input_dims.split(',')
# 收集所有的索引名
index_names = list(set(''.join(input_dims)))
# 收集每个索引名对应的维度大小
sizes = [-1] * len(index_names)
for shape, tensor in zip(input_dims, operands):
for index_name, size in zip(shape, tensor.shape):
assert sizes[index_names.index(index_name)] == -1 or \
sizes[index_names.index(index_name)] == size
sizes[index_names.index(index_name)] = size
# 计算输出矩阵的尺寸
output_size = [sizes[index_names.index(name)] for name in output_dim]
# 将输出矩阵用0初始化
output = operands[0].new_zeros(output_size)
# 遍历所有的索引
for idx in iter_elements(sizes):
# 映射到输出矩阵的索引
idx_output = tuple(idx[index_names.index(name)] for name in output_dim)
prod = 1
for input_dim, tensor in zip(input_dims, operands):
# 对于每一个输入tensor，取得对应的索引
idx_input = tuple(idx[index_names.index(name)] for name in input_dim)
# 计算累积的乘积
prod = prod * tensor[idx_input]
# 求和
output[idx_output] += prod
return output 接下来我们编写一些测试用例，检查einsum_forward和torch.einsum的计算结果是否一致。
# 矩阵乘法
a = torch.rand((5, 4))
b = torch.rand((4, 5))
c = einsum_forward('ij,jk->ik', a, b)
c2 = torch.einsum('ij,jk->ik', a, b)
c3 = a @ b
assert (c - c2).abs().mean() < eps
assert (c - c3).abs().mean() < eps
# 矩阵转置
a = torch.rand((5, 4))
b = einsum_forward('ij->ji', a)
b2 = torch.einsum('ij->ji', a)
b3 = a.T
assert (b - b2).abs().mean() < eps
assert (b - b3).abs().mean() < eps
# 矩阵的迹
a = torch.rand((5, 5))
b = einsum_forward('ii->', a)
b2 = torch.einsum('ii->', a)
b3 = torch.trace(a)
assert (b - b2).abs().mean() < eps
assert (b - b3).abs().mean() < eps
print('全部检查通过。')全部检查通过。4 einsum函数的梯度
举一个简单的例子。假设\(c_{ik} = \sum_{j} a_{ij} b_{jk}\). 对于任意的\(i,j\)，导数\(\frac{\partial L}{\partial c_{ik}}\)都已知，问如何计算\(\frac{\partial L}{\partial a_{ij}}\)和\(\frac{\partial L}{\partial b_{jk}}\)？
显然 \[ \frac{\partial L}{\partial a_{ij}} = \sum_k \frac{\partial L}{\partial c_{ik}}b_{jk}, \] \[ \frac{\partial L}{\partial b_{jk}} = \sum_i \frac{\partial L}{\partial c_{ik}} a_{ij}. \] 可以看到，在这个例子中，不论是前向的计算还是梯度的计算都可以表示为一系列乘积的和。
如果你仔细地观察和理解了einsum函数的计算方式，不难得出更一般的结论，一个由einsum函数定义的表达式，其对任意一个输入参数的梯度同样可以表示为einsum函数的形式。
现在假设a, b, c是某个einsum操作的输入，经过操作我们得到输出d，如下所示：
a = torch.rand((2, 2), requires_grad=True)
b = torch.rand((2, 2), requires_grad=True)
c = torch.rand((2, 2), requires_grad=True)
d = torch.einsum('ij,jk,kl->il', a, b, c)假设\(L\)是\(d\)经过某个函数计算得到的实数，\(\frac{\partial L}{\partial d}\)已知，记为grad_d：
grad_d = torch.rand(d.shape) # 随机初始化grad_d，假设它就是梯度pytorch能很方便地帮我们分别算出a, b, c的梯度。
d.backward(grad_d)
print(a.grad.shape, b.grad.shape, c.grad.shape)torch.Size([2, 2]) torch.Size([2, 2]) torch.Size([2, 2])但是如前所述，einsum函数的梯度仍然可以组织成einsum的形式。让我们来验证这一点。
首先是a的梯度\(\frac{\partial L}{\partial a}\)
grad_a = torch.einsum('il,jk,kl->ij', grad_d, b, c)
assert torch.mean(torch.abs(a.grad - grad_a)) < eps 接著是b的梯度：
grad_b = torch.einsum('il,ij,kl->jk', grad_d, a, c)
assert torch.mean(torch.abs(b.grad - grad_b)) < eps c的梯度：
grad_c = torch.einsum('il,ij,jk->kl', grad_d, a, b)
assert torch.mean(torch.abs(c.grad - grad_c)) < eps 至此，我们验证了a, b, c三个参数的梯度计算方式，它们都可以用einsum函数表示，而且计算结果和pytorch的计算结果一致！
5 总结
最近在实现自己的深度学习模型时，我设计了一个简单的模块，其中一部分用到了einsum。为了加速这个模块在梯度反传阶段的计算，我作了详细的推导，发现einsum有一些有趣的性质，于是就记录下来，形成了这篇文章。
仔细分析可以发现，其实einsum本质上是计算元素间乘积的和的过程。尽管einsum的使用方式多种多样，但只要把握这个本质，其性质也就不难理解了。
