PyTorch tensors topk for every tensor across a dimension

I did it like this:

val, ind = inp[:, :, 2].squeeze().topk(k=4, dim=1, sorted=True)
new_ind = ind.unsqueeze(-1).repeat(1,1,3)
result = inp.gather(1, new_ind)

I don't know if this is the best way to do this but it worked.

One way to do this is by combining fancy indexing and broadcasting as follows:

I am taking a random tensor x of shape (3, 4, 3) and k to be 2 as the example.

>>> import torch
>>> x = torch.rand(3, 4, 3)
>>> x
tensor([[[0.0256, 0.7366, 0.2528],
         [0.5596, 0.9450, 0.5795],
         [0.8265, 0.5469, 0.8304],
         [0.4223, 0.5206, 0.2898]],

        [[0.2159, 0.0369, 0.6869],
         [0.4556, 0.5804, 0.3169],
         [0.8194, 0.5240, 0.0055],
         [0.8357, 0.4162, 0.3740]],

        [[0.3849, 0.0223, 0.9951],
         [0.2872, 0.5952, 0.6570],
         [0.1433, 0.8450, 0.6557],
         [0.0270, 0.9176, 0.3904]]])

Now sort the tensor along the required dimension (here last) and get the indices:

>>> _, idx = torch.sort(x[:, :, -1])
>>> k = 2
>>> idx = idx[:, :k]
# idx is = 
tensor([[0, 3],
        [2, 1],
        [3, 2]])

Now generate three pair of indices (i, j, k) to slice the original tensor as follows:

>>> i = torch.arange(x.shape[0]).reshape(x.shape[0], 1, 1)
>>> j = idx.reshape(x.shape[0], -1, 1)
>>> k = torch.arange(x.shape[2]).reshape(1, 1, x.shape[2])

Note that once you index anything by (i, j, k), they are going to expand and take the shape (x.shape[0], k, x.shape[2]) which is the desired output shape here. Now just index x by i, j and k:

>>> x[i, j, k]
tensor([[[0.0256, 0.7366, 0.2528],
         [0.4223, 0.5206, 0.2898]],

        [[0.8194, 0.5240, 0.0055],
         [0.4556, 0.5804, 0.3169]],

        [[0.0270, 0.9176, 0.3904],
         [0.1433, 0.8450, 0.6557]]])

Essentially, the general recipe that I follow is to create the corresponding access pattern of the tensor via the index arrays and then slicing the tensor directly by using those arrays as indices.

I actually did this for an ascending order sort, so here I am getting top-k least elements. An easy workaround to get the reverse would be to use torch.sort(x[:, :, -1], descending = True).