import torch import numpy as np a = torch.arange(9).reshape(3, 3)提取矩阵对角线元素
out = torch.einsum('ii->i', a) # tensor([0, 4, 8])
矩阵转置
out = torch.einsum('ij->ji', a)
out = torch.einsum('...ij->...ji', a) # 高维矩阵最后两维转置
reduce sum
out = torch.einsum('ij->', a) # tensor(36)
矩阵按列求和
out = torch.einsum('ki->i', a)
矩阵向量乘法
a = torch.arange(6).reshape(2, 3)
b = torch.arange(3)
out = torch.einsum('ik,k->i', a, b)
out = torch.einsum('ik,k', a, b) # 箭头右侧符号可以不写，按规则默认推理。
矩阵乘法
a = torch.arange(6).reshape(2, 3)
b = torch.arange(15).reshape(3, 5)
out = torch.einsum('ik,kj->ij', a, b)
out = torch.einsum('ik,kj', a, b)
向量内积
a = torch.arange(3)
b = torch.arange(3, 6)
out = torch.einsum('i,i->', a, b)
out = torch.einsum('i,i', a, b)
矩阵元素对应相乘并求reduce sum
a = torch.arange(6).reshape(2, 3)
b = torch.arange(6,12).reshape(2, 3)
out = torch.einsum('ij,ij->', a, b)
向量外积
a = torch.arange(3)
b = torch.arange(3,7)
out = torch.einsum('i,j->ij', a, b)
batch矩阵乘法
a = torch.randn(2,3,5)
b = torch.randn(2,5,4)
out = torch.einsum('bik,bkj->bij', a, b)
张量收缩
tensor contraction, 用不上，暂时看不懂。
a = torch.randn(2,3,5,7)
b = torch.randn(11,13,3,17,5)
out = torch.einsum('pqrs,tuqvr->pstuv', a, b)
双线性变换
bilinear transformation. Applies a bilinear transformation to the incoming data.
a = torch.randn(2,3)
b = torch.randn(5,3,7)
c = torch.randn(2,7)
out = torch.einsum('ik,jkl,il->ij', a, b, c)
