【Pytorch】torch.einsumの引数・使い方を徹底解説！アインシュタインの縮約規則を利用して複雑なテンソル操作を短い文字列を使って行う

# 対角成分の和 (Trace)
matrix = torch.randn(4, 4)
trace = torch.einsum('ii', matrix)
print("対角成分の和")
print(trace)  # 行列の対角成分の和（トレース）が出力されます。
print(trace.shape) # 出力: torch.Size([])

# 対角行列
matrix = torch.randn(4, 4)
diag_matrix = torch.einsum('ii->i', matrix)
print("対角行列")
print(diag_matrix)  # 行列の対角成分が出力されます。
print(diag_matrix.shape) # 出力: torch.Size([4])

# ベクトルの作成
x = torch.tensor([[1, 2, 3]])

# 転置
y = torch.einsum('ij->ji', x)

print(y) # 出力: tensor([[1], [2], [3]])

# ベクトルの作成
x = torch.tensor([1, 2, 3])

# 合計
y = torch.einsum('i->', x)

print(y) # 出力: tensor(6)

# 列ごとの合計
matrix = torch.randn(4, 4)
col_sum = torch.einsum('ij->j', matrix)
print("列ごとの合計")
print(col_sum)  # 各列の合計が出力されます。
print(col_sum.shape) # 出力: torch.Size([4])

# 行ごとの合計
matrix = torch.randn(4, 4)
row_sum = torch.einsum('ij->i', matrix)
print("行ごとの合計")
print(row_sum)  # 各行の合計が出力されます。
print(row_sum.shape) # 出力: torch.Size([4])

# 全成分の合計
matrix = torch.randn(4, 4)
total_sum = torch.einsum('ij->', matrix)
print("全成分の合計")
print(total_sum)  # 全成分の合計が出力されます。
print(total_sum.shape) # 出力: torch.Size([])

# ベクトルの作成
x = torch.tensor([1, 2, 3])
z = torch.tensor([4, 5, 6])

# 要素ごとの積
y = torch.einsum('i,i->i', x, z)

print(y) # 出力: tensor([ 4, 10, 18])

# 外積
x = torch.randn(5)
y = torch.randn(4)
outer_product = torch.einsum('i,j->ij', x, y)
print("外積")
print(outer_product)  # ベクトルxとyの外積が出力されます。
print(outer_product.shape) # 出力: torch.Size([5, 4]

# ベクトルの作成
x = torch.tensor([1, 2, 3])
z = torch.tensor([4, 5, 6])

# 外積
y = torch.einsum('i,j->ij', x, z)

print(y) # 出力: tensor([[ 4,  5,  6],
#                        [ 8, 10, 12],
#                        [12, 15, 18]])

# ベクトルの内積
x = torch.randn(5)
y = torch.randn(5)
inner_product = torch.einsum('i,i->', x, y)
print("ベクトルの内積")
print(inner_product)  # ベクトルxとyの内積が出力されます。
print(inner_product.shape) # 出力: torch.Size([])

# ベクトルの作成
x = torch.tensor([1, 2, 3])
z = torch.tensor([4, 5, 6])

# ドット積
y = torch.einsum('i,i->', x, z)

print(y) # 出力: tensor(32)

# 行列とベクトルの作成
x = torch.tensor([[1, 2, 3], [4, 5, 6]])
z = torch.tensor([7, 8, 9])

# 行列-ベクトル積
y = torch.einsum('ij,j->i', x, z)

print(y) # 出力: tensor([ 50, 122])

# 行列の作成
x = torch.tensor([[1, 2, 3], [4, 5, 6]])
z = torch.tensor([[7, 8], [9, 10], [11, 12]])

# 行列-行列積
y = torch.einsum('ik,kj->ij', x, z)

print(y) # 出力: tensor([[ 58,  64],
#                        [139, 154]])

# バッチ行列積
As = torch.randn(3, 2, 5)
Bs = torch.randn(3, 5, 4)
batched_matmul = torch.einsum('bij,bjk->bik', As, Bs)
print("バッチ行列積")
print(batched_matmul)  # バッチ行列積が出力されます。
print(batched_matmul.shape) # 出力: torch.Size([3, 2, 4])

# 3次元テンソルの作成
x = torch.tensor([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]])
z = torch.tensor([[[13, 14], [15, 16], [17, 18]], [[19, 20], [21, 22], [23, 24]]])

# バッチ行列積
y = torch.einsum('bik,bkj->bij', x, z)

print(y) # 出力: tensor([[[ 94, 100],
#                        [229, 244]],
#                        [[508, 532],
#                        [697, 730]]])