torch.einsum函数笔记(爱因斯坦求和约定)

爱因斯坦求和约定

  爱因斯坦求和约定是可以将许多矩阵运算(向量内积、矩阵装置、高维矩阵乘法)简化的一套规则；利用这个函数，可以很方便地实现一些复杂的操作，同时使得代码紧凑；

# 矩阵乘法
a = torch.rand(4,3)
b = torch.rand(3,4)
c = torch.einsum("ik,kj->ij", [a, b])

上式表示将a矩阵与b矩阵相乘重点关注索引的规律：$ik,kj->ij$ ，先看自由标与求和标（哑标）
自由标：在箭头左右都出现的下标；
求和标（哑标）：只在箭头左边出现，不在箭头右侧的下标；
自由标表示运算结果中的排列方式，如上例的箭头右侧的 $ij$ 表示运算的结果是个二维的张量(矩阵)，且矩阵的下标索引为 $ij$；而求和标 $k$ 表示在运算过程中，需要对其维度进行求和，在代码中表示为：

# 遍历输出的坐标，逐个元素求解
for i in range(0, i_max):
	for j in range(0, j_max):
		# 对应矩阵乘法中的，c[i, j] = a[i, 1]*b[1, j] + a[i, 2]*b[2, j] + .... + a[i, k_max]*b[k_max, j]
		tmpSum = 0 
		for k in range(0, k_max):
			tmpSum += (a[i, k]*b[k, j])
		c[i, j] = tmpSum;

特殊情况：

1. 箭头右侧可以不填，此时按照默认规则推导得到结果，把求和索引取出来在内层求和，然后剩余的字母按字母表顺序排列在外层遍历；
2. 支持省略号"…" 表示不关心的索引，若对张量最后两个维度进行转置，可以用 $torch.sinsum{(}^{\prime }...ik->...j{i}^{\prime },a)$