Hi Standshik!



 standshik:

(topk(x))_i =x_i if x_i is one of the top k elements of x, 0 otherwise.


You can do this using torch.topk() and .scatter_():
>>> torch.__version__
'1.7.1'
>>> _ = torch.manual_seed (2021)
>>> x = torch.randperm (10).float()
>>> x.requires_grad = True
>>> x_topk_val_ind = torch.topk (x, 5)
>>> x_topk_val_ind
torch.return_types.topk(
values=tensor([9., 8., 7., 6., 5.], grad_fn=<TopkBackward>),
indices=tensor([8, 6, 3, 5, 1]))
>>> topk = torch.zeros_like (x).scatter_ (0, x_topk_val_ind[1], x_topk_val_ind[0])
>>> topk
tensor([0., 5., 0., 7., 0., 6., 8., 0., 9., 0.], grad_fn=<ScatterBackward0>)
>>> topk.sum().backward()
>>> x.grad
tensor([0., 1., 0., 1., 0., 1., 1., 0., 1., 0.])
>>> x
tensor([4., 5., 3., 7., 2., 6., 8., 0., 9., 1.], requires_grad=True)

Best.
K. Frank