一、前言
einsum 是一个非常强大的函数，用于执行张量（Tensor）运算。它的名称来源于爱因斯坦求和约定（Einstein summation convention），在PyTorch中，einsum 可以方便地进行多维数组的操作和计算。
在Transfomer中，einsum用的非常多，比如使用 einsum 实现自注意力机制中注意力权重的获取，也就是Q和K的内积：
Q（Query）：形状为 (batch_size, seq_len, d_k)
K（Key）：形状为 (batch_size, seq_len, d_k)
import torch
import torch.nn.functional as F
Q = torch.randn(2, 10, 64) # (batch_size, seq_len, d_k)
K = torch.randn(2, 10, 64) # (batch_size, seq_len, d_k)
# (batch_size, seq_len, seq_len)
attention_scores = torch.einsum('bqd,bkd->bqk', Q, K) / torch.sqrt(torch.tensor(64.0))
# (batch_size, seq_len, seq_len)
attention_weights = F.softmax(attention_scores, dim=-1)
二、常见用法示例
2.1 向量点积
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
result = torch.einsum('i,i->', a, b)
print(result) # 输出 32这里，'i,i->' 表示对向量 a 和 b 进行点积操作，其中 i 是索引表示，-> 之后为空表示求和。
2.2 矩阵乘法
A = torch.tensor([[1, 2], [3, 4]])
B = torch.tensor([[5, 6], [7, 8]])
result = torch.einsum('ij,jk->ik', A, B)
print(result) # 输出 tensor([[19, 22], [43, 50]])这里，'ij,jk->ik' 表示矩阵乘法，其中 i 和 k 是结果的维度，j 是求和维度。
2.3 批量矩阵乘法
A = torch.randn(2, 3, 4)
B = torch.randn(2, 4, 5)
result = torch.einsum('bij,bjk->bik', A, B)这里，'bij,bjk->bik' 表示对批量的矩阵进行乘法运算。
解释：
bij，bjk分别是A和B的3个维度，用字符串的形式指代。
为什么最后得到的是bik呢？这个和线性代数的矩阵运算规则有关系。
矩阵乘法规则：
给定矩阵 A 的形状为 (m,n)
给定矩阵 B 的形状为 (n,p)
矩阵乘法 A×B 的结果矩阵 C 的形状为 (m,p)
在矩阵乘法中，结果矩阵的每个元素 Cik 是通过 A 的第 i 行和 B 的第 k 列的对应元素相乘并求和得到的，即：
计算过程：
1. 匹配批次维度 (b)
对于每个批次，独立进行矩阵乘法运算。
2. 求和维度 (j)：
j 是两个张量中共同的维度，根据线性代数中的矩阵乘法规则，需要对 j 维度进行求和。
3. 保留和产生的维度：
i 来自 A，表示保留 A 的第一个维度。
k 来自 B，表示保留 B 的第二个维度。
经过上述分析，einsum 的结果保留了 b（批次维度）、i（来自 A 的第一个维度）和 k（来自 B 的第二个维度）。因此，结果张量的形状为 (batch_size, seq_len_i, seq_len_k)，也就是 bik。
同样，延伸到4维计算的话。
torch.einsum("nqhd,nkhd->nhqk", [queries, keys])首先，假设 queries 和 keys 的形状为：
queries: (batch_size, seq_len_q, num_heads, head_dim)
keys: (batch_size, seq_len_k, num_heads, head_dim)
用具体变量名表示：
n: batch_size，批次大小。
q: seq_len_q，查询序列的长度。
k: seq_len_k，键序列的长度。
h: num_heads，多头注意力中的头数。
d: head_dim，每个头的维度。
1. 匹配批次维度 (n) 和头部维度 (h)：
批次大小和头部数量在两个输入张量中都是相同的，保持不变。
2. 求和维度 (d)：
d 表示每个头的维度。在 queries 和 keys 中，d 都是最后一个维度，对这个维度进行点积运算后求和。
3. 保留和产生的维度：
q 来自 queries，表示查询序列的长度。
k 来自 keys，表示键序列的长度。
所以最后是nhqk。
2.4 转置操作
A = torch.tensor([[1, 2, 3], [4, 5, 6]])
result = torch.einsum('ij->ji', A)
print(result) # 输出 tensor([[1, 4], [2, 5], [3, 6]])这里，'ij->ji' 表示将矩阵进行转置操作。