关于torch.einsum函数‘ii’模式输出结果的疑问

示例代码 #

A = torch.randn(5,5)
B = torch.einsum("ii->i",A)
C = torch.einsum("ii",A)

运行结果 #

before: tensor([[-0.2339,  0.2501, -1.1814,  1.4392, -0.5461],
        [ 1.4908,  0.0626, -0.6849, -1.3106,  0.1257],
        [ 3.3362, -1.7438,  0.3027,  0.4346,  0.6830],
        [-0.6183,  0.5965,  1.2653,  1.0319, -0.0670],
        [-0.5531, -0.4245, -2.4869,  1.2972,  0.6732]])
after: tensor([-0.2339,  0.0626,  0.3027,  1.0319,  0.6732])
Why: tensor(1.8366)

结果解释 #

• torch.einsum("ii->i", A)：箭头->明确指定了输出维度为i，因此会提取张量A的对角线元素，输出一个一维张量，对应结果里的after。
• torch.einsum("ii", A)：当没有箭头指定输出维度时，einsum会对重复的索引执行求和操作。这里重复索引是i，所以会把对角线上的所有元素相加：

  -0.2339 + 0.0626 + 0.3027 + 1.0319 + 0.6732 = 1.8366
  

  最终得到一个标量，对应结果里的Why。

内容的提问来源于stack exchange，提问作者刘炜杰