kth-value per row in pytorch?

user_7034613 0
Viewed 479

Given 

import torch    
A = torch.rand(9).view((3,3)) # tensor([[0.7455, 0.7736, 0.1772],\n[0.6646, 0.4191, 0.6602],\n[0.0818, 0.8079, 0.6424]])
k = torch.tensor([0,1,0])
A.kthvalue_vectoriezed(k) -> [0.1772,0.6602,0.0818]

Meaning I would like to operate on each column with a different k.
Not kthvalue nor topk offers such API. Is there a vectorized way around that?
Remark - kth value is not the value in the kth index, but the kth smallest element. Pytorch docs

torch.kthvalue(input, k, dim=None, keepdim=False, out=None) -> (Tensor, LongTensor)
Returns a namedtuple (values, indices) where values is the k th smallest element of each row of the input tensor in the given dimension dim. And indices is the index location of each element found.

pytorch vectorization
1 Answers

Assuming you don't need indices into original matrix (if you do, just use fancy indexing for the second return value as well) you could simply sort the values (by last index by default) and return appropriate values like so:

def kth_smallest(tensor, indices):
    tensor_sorted, _ = torch.sort(tensor)
    return tensor_sorted[torch.arange(len(indices)), indices]

And this test case gives you your desired values:

tensor = torch.tensor(
    [[0.7455, 0.7736, 0.1772], [0.6646, 0.4191, 0.6602], [0.0818, 0.8079, 0.6424]]
)

print(kth_smallest(tensor, [0, 1, 0])) # -> [0.1772,0.6602,0.0818]
user_108864200
Related