torch topk

torch.topk 是 PyTorch 中的一个函数，用于返回输入张量中 k 个最大值或最小值及其对应的索引。它的语法如下：

torch.topk(input, k, dim=None, largest=True, sorted=True, out=None) -> (Tensor, LongTensor)

其中，参数含义如下：

- `input`：输入张量。
- `k`：需要返回的最大或最小值的个数。
- `dim`：沿着哪个维度计算，默认为输入张量的最后一维。
- `largest`：为 True 时返回最大值，为 False 时返回最小值，默认为 True。
- `sorted`：为 True 时返回排序后的结果，为 False 时返回未排序的结果，默认为 True。
- `out`：输出张量。

下面是一个示例：

import torch

x = torch.tensor([[1, 3, 2], [4, 2, 1]])
values, indices = torch.topk(x, k=2)
print(values)   # tensor([[3, 2],
                #         [4, 2]])
print(indices)  # tensor([[1, 2],
                #         [0, 1]])

这里输入张量 x 的形状是 (2, 3)，即两行三列。我们指定 k=2，表示需要返回每行中的两个最大值及其对应的索引。因此，输出的 values 张量的形状是 (2, 2)，indices 张量的形状也是 (2, 2)。