Is topk broken?

I just notice that when specifying sorted=False , the values are still sorted.
Have I misunderstood the documentation?

x = torch.rand(5,5)
>>> x
tensor([[0.6847, 0.0039, 0.1289, 0.9117, 0.6705],
        [0.4684, 0.0572, 0.3332, 0.7579, 0.8104],
        [0.9675, 0.7871, 0.9898, 0.3934, 0.4836],
        [0.0480, 0.1028, 0.0757, 0.7755, 0.7489],
        [0.1737, 0.0918, 0.9962, 0.4362, 0.7884]])

>>> x.topk(2, sorted=False).values
tensor([[0.9117, 0.6847],
        [0.8104, 0.7579],
        [0.9898, 0.9675],
        [0.7755, 0.7489],
        [0.9962, 0.7884]])

sorted=True will make sure that the results are indeed sorted. Setting it to False does not “unsort” them, but doesn’t guarantee they will be sorted (they still could be depending on the backend implementation).

Thanks for your answer!

It is certainly a surprise though. On both cpu and cuda the sorted argument does’nt change the result at all.

It does change the output on my setup:

x = torch.rand(5,5).cuda()
print(torch.topk(x, 5, sorted=False).indices == torch.topk(x, 5, sorted=True).indices)
> tensor([[False, False, False, False,  True],
          [ True, False, False, False,  True],
          [ True, False, False,  True,  True],
          [False, False, False, False,  True],
          [False, False, False,  True,  True]], device='cuda:0')

So, it seems to depend on k:

>>> print(torch.topk(x, 2, sorted=False).indices == torch.topk(x, 2, sorted=True).indices)
tensor([[True, True],
        [True, True],
        [True, True],
        [True, True],
        [True, True]], device='cuda:0')