einsum is all you need

最近发现一个超赞的函数：einsum！

每次看到它，我都会想起一句话：

优雅永不过时

函数的名字也很好记：爱因斯坦求和 = einsum（）

注意：矩阵的维度只能使用26个小写英文字母

功能

特殊用法

省略箭头及以后，就是按默认运算，这个尽量慎用，可能会加大看懂代码的工作量。。。

a = torch.randn(1,2)
b = torch.randn(2,3)
c = torch.randn(3,4)
d = torch.einsum('xy,yz,zj',[a,b,c])
print(d.shape)

...表示不需要进行关注的维度

a = torch.randn(1,2,3)
b = torch.einsum('...ij->...ji', [a])

 

1、矩阵乘法

a = torch.randn(1,2)
b = torch.randn(2,3)
c = torch.einsum("ik,kj->ij", [a, b])

并且可以做到多个矩阵连乘

a = torch.randn(1,2)
b = torch.randn(2,3)
c = torch.randn(3,4)
d = torch.einsum('xy,yz,zj->xj',[a,b,c])
print(d.shape)

优雅，太优雅了，这样来能极大地避免矩阵维度的bug，简洁直观一目了然。

还可以把转置也融合进来

a = torch.randn(1,2)
b = torch.randn(2,3)
c = torch.randn(3,4)
d = torch.einsum('xy,yz,zj->jx',[a,b,c])
print(d.shape)

2、提取对角线元素

a = torch.randn(3,3)
b = torch.einsum('ii->i', [a])

3、维度变换

a = torch.randn(1,2,3,4,5)
b = torch.einsum('abcde->edcba', [a])

4、求和

a = torch.randn(3,3)
b = torch.einsum('ij->i', [a]) # 逐列相加
c = torch.einsum('ij->j', [a]) # 逐行相加
d = torch.einsum('ij->', [a]) # 全部相加

5、向量内积

a = torch.arange(3)
b = torch.arange(3)
c = torch.einsum('i,i->', [a, b])

 

6、矩阵对应元素相乘并求和

a = torch.arange(6).reshape(2, 3)
b = torch.arange(6,12).reshape(2, 3)
c = torch.einsum('ij,ij->', [a, b]).numpy()

7、向量外积

a = torch.arange(3)
b = torch.arange(4)
c = torch.einsum('i,j->ij', [a, b])