PyTorch torch.einsum 関数
Pytorch torch リファレンスマニュアルtorch.einsumこれは、PyTorch でアインシュタインの縮約記法に使用される関数で、さまざまなテンソル演算を簡潔に表すことができます。
関数定義
torch.einsum(equation, *operands)
使用例
例
import torch
# 行列転置
A = torch.randn(2, 3)
AT = torch.einsum('ij->ji', A)
print("転置形状:", AT.shape)
# 行列乗算
B = torch.randn(3, 4)
C = torch.einsum('ij,jk->ik', A, B)
print("乗算形状:", C.shape)
# 内積
a = torch.randn(3)
b = torch.randn(3)
dot = torch.einsum('i,i->', a, b)
print("内積:", dot.item())
# 行列転置
A = torch.randn(2, 3)
AT = torch.einsum('ij->ji', A)
print("転置形状:", AT.shape)
# 行列乗算
B = torch.randn(3, 4)
C = torch.einsum('ij,jk->ik', A, B)
print("乗算形状:", C.shape)
# 内積
a = torch.randn(3)
b = torch.randn(3)
dot = torch.einsum('i,i->', a, b)
print("内積:", dot.item())
その他の拡張