torch gather nd

torch.gather_nd 是一个 PyTorch 的函数，用于按照索引从输入张量中收集元素并构建输出张量。它的参数包括：

- input: 输入张量；
- dim: 要在其上索引的维度；
- index: 索引张量，其形状必须与输出张量的形状相同；
- out=None: 可选输出张量。

具体而言，对于输入张量 input 和索引张量 index，输出张量的形状为 index.shape[:-1] + input.shape[index.shape[-1]:]。例如，如果 index.shape 为 (2, 3, 4)，而 input.shape 为 (5, 6, 7, 8)，则输出张量的形状将为 (2, 3, 4, 8)。

该函数的作用是，对于输入张量 input 中的每个元素，按照索引张量 index 中对应的索引从 input 中取出一个元素，并将这些元素组成一个新的张量作为输出。具体来说，如果输出张量的形状为 (i,j,k,l)，则输出张量中的元素 out[i,j,k,l] 的值为 input[index[i,j,k],index[i,j,k],...,index[i,j,k],l]。其中，index[i,j,k] 是一个形状为 (dim,) 的张量，表示对于输入张量 input 的第 dim 维，要从中取出的元素的索引。

下面是一个简单的示例：

import torch

input = torch.tensor([[1, 2], [3, 4], [5, 6]])
index = torch.tensor([[0, 0], [2, 1]])

output = torch.gather_nd(input, dim=0, index=index)
print(output)

该代码输出：

tensor([[1, 2],
        [5, 4]])

其中，output 中的第一个元素为 input[0,0]，第二个元素为 input[2,1]。