torch.gather 函数的作用是按照指定的维度 dim 和索引 index 从输入张量 input 中收集数值。这个操作通常用于根据索引从一个维度中选择元素，并生成一个新的张量作为输出
1. 介绍
1.1 参数说明
- input: 需要从中选取元素的原始张量。
- dim: 沿着此维度选取元素。例如，如果 dim=0，则沿着第一个维度（通常是
行）选取；如果 dim=1，则沿着第二个维度（通常是列）选取。 - index: 一个长整型张量，包含要选取的索引。index 的形状应该与 input 的形状相同，或者可以广播到 input 的形状。
1.2. 索引张量 index 的作用
- index 张量中的每个元素指定了在 input 张量中 dim 维度上的位置。例如，如果
dim=1(列) 并且index[i, j]的值为 k，则从第 i 行的第 k 列选取元素。 - 根据 index 张量中的索引，在 input 张量中沿着 dim 维度收集元素。
- 输出张量的形状与 index 张量的形状相同。这意味着除了 dim 维度之外，其他所有维度的大小都与 index 相同。
2. 示例
import torch# 创建一个输入张量
input_tensor = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])# 创建一个索引张量，其形状与输入张量相同
index_tensor = torch.tensor([[0, 2, 1], [2, 0, 1], [1, 0, 2]])# 使用 torch.gather 收集元素，沿着列（dim=1）
output_tensor = torch.gather(input_tensor, 1, index_tensor)print(output_tensor)
说明
在上面的示例中，torch.gather(input_tensor, 1, index_tensor) 的输出将是：
- 对于第 0 行，列索引(dim为1)为 [0, 2, 1]，所以收集的元素是 [1, 3, 2]。
- 对于第 1 行，列索引为 [2, 0, 1]，所以收集的元素是 [6, 4, 5]。
- 对于第 2 行，列索引为 [1, 0, 2]，所以收集的元素是 [8, 7, 9]。
因此，输出张量将是：
tensor([[1, 3, 2],[6, 4, 5],[8, 7, 9]])
注意事项：确保 index 中的所有值都在有效范围内，即从 0 到 input.size(dim) - 1。如果 index 中有任何值超出了这个范围，将会引发错误。
)