torch.kthvalue

torch.kthvalue(input, k, dim=None, keepdim=False, *, out=None)

返回一個命名元組 (values, indices)，其中 values 是給定維度 dim 中 input 張量每行的第 k 小元素。而 indices 是找到的每個元素的索引位置。

如果未給定 dim，則選擇 input 的最後一個維度。

如果 keepdim 為 True，則 values 和 indices 張量的大小與 input 相同，除了維度 dim 的大小為 1。否則，dim 被壓縮（參見 torch.squeeze()），導致 values 和 indices 張量比 input 張量少 1 個維度。

注意

當 input 是 CUDA 張量且存在多個有效的第 k 小值時，此函式可能不確定性地返回其中任意一個的 indices。

引數

• input (Tensor) – 輸入張量。

• k (int) – 第 k 小元素的 k 值

• dim (int, 可選) – 沿哪個維度尋找第 k 小值

• keepdim (bool) – 輸出張量是否保留 dim。

關鍵字引數

out (tuple, 可選) – 可選地提供 (Tensor, LongTensor) 組成的輸出元組，用作輸出緩衝區

示例

>>> x = torch.arange(1., 6.)
>>> x
tensor([ 1.,  2.,  3.,  4.,  5.])
>>> torch.kthvalue(x, 4)
torch.return_types.kthvalue(values=tensor(4.), indices=tensor(3))

>>> x=torch.arange(1.,7.).resize_(2,3)
>>> x
tensor([[ 1.,  2.,  3.],
        [ 4.,  5.,  6.]])
>>> torch.kthvalue(x, 2, 0, True)
torch.return_types.kthvalue(values=tensor([[4., 5., 6.]]), indices=tensor([[1, 1, 1]]))