Topk on 2D dimension

1

I am trying to run a topk of a 2d matrix

x = torch.tensor([[1, 2,9,87],[6, 32,8,1],[4,6,7,2],[3,6,2,6]])
print("x",x)
values, indices = torch.topk(x,k=2,dim=0)
print("values",values)
print("indices",indices)

output:

x tensor([[ 1,  2,  9, 87],
        [ 6, 32,  8,  1],
        [ 4,  6,  7,  2],
        [ 3,  6,  2,  6]])
values tensor([[ 6, 32,  9, 87],
        [ 4,  6,  8,  6]])
indices tensor([[1, 1, 0, 0],
        [2, 2, 1, 3]])

but I am expecting output for values as:

values tensor([[ 6, 32,  8,  1],
        [ 4,  6,  7,  2]])
2

The topk works column by column, and you get the top values there.
From what you write, do you expect the rows with the largest leading values? I’m afraid there is an easy way to get those…
Depending on what you know about your values, just running topk on the first column and using those indices on x might work.

Best regards

Thomas

3

Thank you… I ended up doing the following:
I needed the last 3 elements of each array so using ranges

x = torch.tensor([[1, 0,9,87,1],[6, 0,8,1,8],[4,0,7,2,2],[3,0,2,6,3]])
print("x",x)
values, indices = torch.topk(x,k=2,largest=True,dim=0)
print("indices",indices)
#2d data samples
newdataset = torch.zeros(2, 3)
print("newdataset",newdataset)
newdataset[0] = x[indices[0,0].item(),2:5]
newdataset[1] = x[indices[1,0].item(),2:5]
print("newdataset",newdataset)
4

For what it’s worth, a slightly more concise way to do this (that also avoids creating a new tensor, which I find can be cumbersome for autograd purposes) would be

>>> x = torch.tensor([
        [1, 0, 9, 87, 1],
        [6, 0, 8, 1, 8],
        [4, 0, 7, 2, 2],
        [3, 0, 2, 6, 3]
    ])
>>> _, indices = torch.topk(x[:, 0], 2)
>>> indices
tensor([1, 2])
>>> new_dataset = x[indices, 2:5]
>>> new_dataset
tensor([[8., 1., 8.],
        [7., 2., 2.]])