引言
我想向您介绍Python中最有用的方法之一，np.einsum。
使用np.einsum（以及Tensorflow和JAX中的对应方法），您可以以极其清晰和简洁的方式编写复杂的矩阵和张量操作。我还发现，它的清晰性和简洁性减轻了处理张量时所带来的大量心理负担。
而且，它实际上相对简单易学和使用。它的工作原理如下：
在np.einsum中，您有一个subscripts字符串参数和一个或多个操作数：
numpy.einsum(subscripts : string, *operands : List[np.ndarray])
参数 subscripts 是一种“迷你语言”，用于告诉 numpy 如何操作和组合操作数的轴。起初可能有些难以理解，但掌握后就会觉得并不复杂。
单一操作数
作为第一个例子，我们使用 np.einsum 来交换矩阵 A 的轴（即进行转置）：
M = np.einsum('ij->ji', A)
字母 i 和 j 绑定到 A 的第一和第二个轴。Numpy 按照字母出现的顺序将字母绑定到轴，但如果你明确指定，Numpy 不在乎你使用什么字母。例如，我们可以使用 a 和 b，效果是一样的：
M = np.einsum('ab->ba', A)
然而，您必须提供与操作数中的轴数量相同的字母。A中有两个轴，因此您必须提供两个不同的字母。下一个示例将不会工作，因为下标公式只有一个字母可绑定，即i:
# broken
M = np.einsum('i->i', A)
另一方面，如果操作数确实只有一个轴（换句话说，它是一个向量），那么单字母下标公式也能正常工作，尽管它并不是很有用，因为它保持了向量a不变:
m = np.einsum('i->i', a)
在轴上求和
但是这个操作呢？右侧没有 i。这有效吗？
c = np.einsum('i->', a)
令人惊讶的是，是的！理解np.einsum本质的第一个关键是：如果右侧的一个轴被省略，那么该轴将被求和。
代码：
c = 0
I = len(a)
for i in range(I):
c += a[i]
求和操作并不限于单一轴。 例如，您可以通过使用这个下标公式同时对两个轴进行求和：c = np.einsum('ij->', A)：
以下是相应的 Python 代码：
c = 0
I,J = A.shape
for i in range(I):
for j in range(J):
```python
c += A[i,j]
```
但是这并不是全部——我们可以发挥创意，对某些轴进行求和，而对其他轴保持不变。例如：`np.einsum('ij->i', A)` 对矩阵 `A` 的行进行求和，留下一个长度为 `j` 的行和向量：
代码:
I,J = A.shape
r = np.zeros(I)
for i in range(I):
for j in range(J):
r[i] += A[i,j]
同样，np.einsum('ij->j', A) 对 A 的列进行求和。
代码：
I,J = A.shape
r = np.zeros(J)
for i in range(I):
for j in range(J):
r[j] += A[i,j]
两个操作数
使用单个操作数的能力是有限的。使用两个操作数时，事情变得更加有趣（也更有用）。
假设你有两个向量 a = [a_1, a_2, ... ] 和 b = [a_1, a_2, ...]。
如果 len(a) === len(b)，我们可以这样计算 内积（也称为点积）：
a = np.asarray([4,5,6])
b = np.asarray([1,2,3])
c = np.einsum('i,i->', a, b)`
>> c := 32.0
这里同时发生了两件事情：
- 因为
i同时绑定于a和b，所以a和b被“对齐”，然后相乘：a[i] * b[i]。 - 由于索引
i被排除在右侧，因此在轴i上进行求和以消除它。
将 (1) 和 (2) 结合在一起，你就得到了经典的内积。
代码：
c = 0
I = len(a)
for i in range(I):
c += a[i] * b[i]
现在，假设我们没有在下标公式中省略i，我们将对所有a[i]和b[i]进行相乘，而不是对i求和：
a = np.asarray([4,5,6])
b = np.asarray([1,2,3])
c = np.einsum(`i,i->i`, a, b)
>> c := np.asarray([4,10,18])
代码:
I = len(a)
c = np.zeros(I)
for i in range(I):
c[i] = a[i] * b[i]
这也被称为逐元素乘法（对于矩阵来说是哈达玛乘积），通常通过numpy方法 np.multiply 来实现。
最后，假设我们在输出中包含了 所有 的轴 – 即 i 和 j。这被称为 外积。
a = np.asarray([4,5,6])
b = np.asarray([1,2,3])
C = np.einsum(`i,j->ij`, a, b)
>> C := np.asarray([[4,8,12],[5,10,15],[6,12,18]])
在这个下标公式中，a 和 b 的轴被绑定到不同的字母，因此被视为独立的“循环变量”。因此，C 的条目为 a[i] * b[j]，对于所有的 i 和 j，这些条目被排列成一个矩阵。
代码:
I = len(a)
J = len(b)
C = np.zeros(I,J)
for i in range(I):
for j in range(J):
C[i,j] = a[i] * b[j]
三个操作数
将外积进一步扩展，这里是一个三个操作数的版本：
M = np.einsum('i,j,k->ijk', a, b, c)
我们三操作数外积的等效Python代码为：
I = len(a)
J = len(b)
K = len(c)
for i in range(I):
for j in range(J):
for k in range(K):
M[i,j,k] = a[i] * b[j] * c[k]
更进一步，我们可以选择省略某些轴以对它们进行求和，同时通过在->的右侧写ki而不是ik来转置结果：
M = np.einsum('i,j,k->ki', a, b, c)
等效的Python代码如下：
I = len(a)
J = len(b)
K = len(c)
M = np.zeros(K,I)
for i in range(I):
for j in range(J):
for k in range(K):
M[k,i] += a[i] * b[j] * c[k]
现在我希望你能开始看到如何相对容易地指定复杂的张量操作。更重要的是，我可以直接从下标中读取上述操作：“三个向量的外积，中间轴求和，最终结果转置”。这很不错，但这仅仅是学术上的吗？我不这样认为。
一个实际的例子
作为一个实际的例子，让我们实现LLMs核心的方程，来自经典论文“Attention is All You Need”。
公式1描述了注意力机制：
我们将重点关注术语
由于 softmax 不能通过 np.einsum 计算，并且缩放因子表示 m 个查询与 n 个键的点积。Q 是一个包含 m 个 d 维行向量的集合，这些行向量堆叠成一个矩阵，因此 Q 的形状为 md。同样，K 是一个包含 n 个 d 维行向量的集合，这些行向量也堆叠成一个矩阵，因此 K 的形状为 md。
单个 Q 和 K 之间的乘积可以写作：
np.einsum('md,nd->mn', Q, K)
请注意，由于我们编写下标方程的方式，我们避免了在矩阵乘法之前转置 K！
所以，这看起来相当简单——实际上，这只是传统的矩阵乘法。然而，我们还没有完成。Attention Is All You Need 使用了 多头注意力，这意味着我们实际上有 k 个这样的矩阵乘法同时发生在一个索引集合的 Q 矩阵和 K 矩阵上。
为了使事情更清晰，我们可以将乘积重写,
这意味着我们为 Q 和 K 增加了一个额外的轴 i。
更重要的是，如果我们处于训练环境中，我们可能正在执行一批这样的多头注意力操作。
因此，我们可能希望沿着批次轴 b 对一批示例执行该操作。因此，完整的计算可以表示为：
batch_multihead_QKt = np.einsum('bimd,bind->bimn', Q, K, optimize = True)
我将跳过这里的图示，因为我们正在处理4轴张量。但你可能可以想象将之前的图示“堆叠”起来以获得我们的多头轴 i，然后再将这些“堆叠”起来以获得我们的批次轴 b。
我很难想象我们如何使用其他numpy方法的任何组合来实现这样的操作。然而，通过一点检查，很明显发生了什么：在一个批次上，对一组矩阵Q和K，执行矩阵乘法Qt(K)。
现在，这不是很棒吗？
