torch.einsum函数的使用

torch.einsum函数可以实现多种张量运算，如矩阵乘法、向量点积、外积、迹、对角化等。

基本语法：

torch.einsum(equation, *operands)

• equation 是一个描述运算的字符串，格式为 'inputs->outputs'，其中 inputs 描述输入张量的维度，outputs 描述输出张量的维度。
• *operands 是一个或多个输入张量。

基本用法

• 矩阵乘法

假设有两个矩阵 (A) 和 (B)，维度分别为 (m x n) 和 (n x p)。矩阵乘法的 einsum 表示为：

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

 ij 代表矩阵 (A) 的两个维度，jk 代表矩阵 (B) 的两个维度，ik 代表输出矩阵 (C) 的维度。

• 向量点积

对于两个向量 (a) 和 (b)，其点积可以表示为：

dot_product = torch.einsum('i,i->', a, b)

• 外积

对于两个向量 (a) 和 (b)，其外积（即生成一个矩阵，其中每个元素是 ）为：

outer_product = torch.einsum('i,j->ij', a, b)

• 迹

对于一个矩阵 (A)，其迹（即对角线元素之和）为：

trace = torch.einsum('ii->', A)

• 对角化

将一个向量 (v) 转换为一个对角矩阵，其对角线元素为 (v) 的元素：

diag_matrix = torch.einsum('i,j->ij', v, torch.arange(len(v)) == torch.arange(len(v))[:, None])

diag_matrix = torch.einsum('i,j->ij', v, (torch.eye(len(v)) == 1).float())
#这里使用 (torch.eye(len(v)) == 1).float() 是为了生成一个对角线上全为 1，其余位置为 0 的浮点矩阵。

• 批量矩阵乘法

假设你有两个三维张量 (A) 和 (B)，维度分别为 (b x m x n) 和 (b x n x p)，其中 (b) 是批量大小。批量矩阵乘法的 einsum 表示为：

C = torch.einsum('bij,bjk->bik', A, B)

• 多矩阵乘法

假设有三个矩阵 (A), (B), 和 (C)，维度分别为 (l x m), (m x n), 和 (n x o)，计算 (ABC): 

D = torch.einsum('ij,jk,kl->il', A, B, C)