torch.argsort

torch.argsort(input, dim=-1, descending=False, stable=False) → Tensor

傳回沿給定維度以數值升序排列張量的索引。

這是 torch.sort() 傳回的第二個值。有關此方法的確切語義，請參閱其文件。

如果 stable 為 True，則排序常式會變成穩定的，從而保留相等元素的順序。 如果 False，則不保證比較相等的數值的相對順序。True 速度較慢。

參數

• input (Tensor) – 輸入張量。

• dim (int, optional) – 要排序的維度

• descending (bool, optional) – 控制排序順序（遞增或遞減）

• stable (bool, optional) – 控制相等元素的相對順序

範例

>>> a = torch.randn(4, 4)
>>> a
tensor([[ 0.0785,  1.5267, -0.8521,  0.4065],
        [ 0.1598,  0.0788, -0.0745, -1.2700],
        [ 1.2208,  1.0722, -0.7064,  1.2564],
        [ 0.0669, -0.2318, -0.8229, -0.9280]])


>>> torch.argsort(a, dim=1)
tensor([[2, 0, 3, 1],
        [3, 2, 1, 0],
        [2, 1, 0, 3],
        [3, 2, 1, 0]])