两个tensor作为另一个tensor的索引

形如 tensor[tensor, tensor] 的理解如下：

首先看一段代码：

import torch
S=torch.tensor(range(60)).reshape(3,4,5)
print(S)
A=[[0,1,2],
   [2,1,2],
   [0,1,2]]

B=[[1,0,2],
   [2,1,3],
   [0,1,3]]

A=torch.tensor(A)
B=torch.tensor(B)

c=S[A,B]
print(c)
print(c.shape)

程序的输出如下：

tensor([[[ 0,  1,  2,  3,  4],
         [ 5,  6,  7,  8,  9],
         [10, 11, 12, 13, 14],
         [15, 16, 17, 18, 19]],

        [[20, 21, 22, 23, 24],
         [25, 26, 27, 28, 29],
         [30, 31, 32, 33, 34],
         [35, 36, 37, 38, 39]],

        [[40, 41, 42, 43, 44],
         [45, 46, 47, 48, 49],
         [50, 51, 52, 53, 54],
         [55, 56, 57, 58, 59]]])
tensor([[[ 5,  6,  7,  8,  9],
         [20, 21, 22, 23, 24],
         [50, 51, 52, 53, 54]],

        [[50, 51, 52, 53, 54],
         [25, 26, 27, 28, 29],
         [55, 56, 57, 58, 59]],

        [[ 0,  1,  2,  3,  4],
         [25, 26, 27, 28, 29],
         [55, 56, 57, 58, 59]]])
torch.Size([3, 3, 5])

以下就根据上面的代码来分析，下面的A，B，S等均为代码中的变量。

        首先来看一下结果的维度问题。A，B作为索引不会改变S的最后一维，因此结果的最后一维与原tensor相同，在上述代码中就是5. 结果的前面的维度与A、B中较大的保持一致（我是这么理解的，欢迎大佬指正），然后A，B中较小的要广播到与较大的一致。在上面代码中A、B均为3*3，所以结果c的维度就是3*3*5，且不需要广播。

        下面来看一下具体的计算过程。代码中S维度为[3, 4, 5]，A用来索引第0维，B用来索引第1维，且A与B中的元素一一对应。例如取A[0, 0] 和 B[0, 0] 的值组成索引[0, 1]，用这个索引到S中找到[5, 6, 7, 8, 9]；取A[0, 1] 和 B[0, 1] 的值组成索引[1, 0]，到S中找到[20, 21, 22, 23, 24]。以此类推就可以了。