张量矩阵乘法返回向量 einsum

Question

我对以下 returns 向量的矩阵张量乘法示例感到困惑。乍一看，我认为这意味着将张量 dydx 的第一维乘以矩阵 dLdy，但我没有得到如下所示的预期结果。那么这个einsum是什么意思？

import torch
import numpy as np

dLdy = torch.randn(2,2)
dydx = torch.randn(2,2,2)

torch.einsum('jk,jki->i', dLdy, dydx)
tensor([0.3115, 3.7255])

dLdy
tensor([[-0.4845,  0.6838],
    [-1.1723,  1.4914]])

dydx
tensor([[[ 1.5496, -1.2722],
     [ 0.1221,  1.0495]],

    [[-1.4882,  0.0307],
     [-0.5134,  1.6276]]])

(dLdy * dydx[0]).sum()
-0.1985

Answer 1

对于 A 和 B，这是前两个维度 jk 的收缩（和），所以 res(i) = sum_{j,k} A(j,k)B(j,k,i)

例如：

import torch
import numpy as np


dLdy = torch.randn(2,2)
dydx = torch.randn(2,2,2)

print(torch.einsum('jk,jki->i', dLdy, dydx))
print((dLdy * dydx[:,:,0]).sum())
print((dLdy * dydx[:,:,1]).sum())

产生

tensor([4.6025, 1.8987])
tensor(4.6025)
tensor(1.8987)

即(dLdy * dydx[:,:,0]).sum()是结果向量的第一个元素，等等

张量矩阵乘法返回向量 einsum

Tensor matrix multiplication returning vector einsum

python

pytorch