torch suppress to kth largest values
I have the following function which works, but just not for half precision values (get a NotImplemented error for kthvalue).
def suppress_small_probabilities(probabilities: torch.FloatTensor, k: int) -&gt; torch.FloatTensor:
    kth_largest, _ = (-probabilities).kthvalue(k, dim=-1, keepdim=True)
    return probabilities * (probabilities &gt;= -kth_largest)

How would you do the equivalent without using kthvalue? I'm guessing topk has something to do with it, but I want to suppress the smaller values. probabilities is of size batch_size x 1000.

Implement your own topk, e.g.
def mytopk(xs: Tensor, k: int) -&gt; Tensor:
    mask = torch.zeros_like(xs)
    batch_idx = torch.arange(0, len(xs))
    for _ in range(k):
        _, index = torch.where(mask == 0, xs, -1e4).max(-1)
        mask[(batch_idx, index)] = 1
    return mask

This will return a boolean mask tensor where the row-wise top-k elements will have value 1, rest 0.
Then use the mask to index your original tensor, e.g.
xs = torch.rand(3, 5, dtype=torch.float16)
# tensor([[0.0626, 0.9620, 0.5596, 0.4423, 0.1932],
#        [0.5289, 0.0857, 0.7802, 0.7730, 0.4807],
#        [0.8272, 0.5016, 0.1169, 0.4372, 0.1843]], dtype=torch.float16)
mask = mytopk(xs, 2)
# tensor([[0., 1., 1., 0., 0.],
#        [0., 0., 1., 1., 0.],
#        [1., 1., 0., 0., 0.]])
top_only = torch.where(mask == 1, xs, 0)
# tensor([[0.0000, 0.9620, 0.5596, 0.0000, 0.0000],
#        [0.0000, 0.0000, 0.7802, 0.7730, 0.0000],
#        [0.8271, 0.5016, 0.0000, 0.0000, 0.0000]], dtype=torch.float16)

