The documentation 1 of topk says it returns the k largest elements but I only need the max values and don’t know how many of them will be there.
I tried using max but it results in only the first max value having the gradient, it’s not shared among the max values. Here is an example,
inp = torch.tensor([[ 6.,  2,  3,  6],
                    [ 9, 16, 16, 16]], requires_grad=True)
outp, inds = torch.max(inp, dim=-1)
outp.sum().backward()
print('outp:\n', outp)
print('inp.grad:\n', inp.grad)

outp:
 tensor([ 6., 16.], grad_fn=<MaxBackward0>)
inp.grad:
 tensor([[1., 0., 0., 0.],
        [0., 1., 0., 0.]])

I expect the gradient to be like below,
inp.grad:
 tensor([[0.5, 0., 0., 0.5],
        [0., 0.333, 0.333, 0.333]])

So, I think I have to implement a custom max but I don’t know how to reshape the input (or squeeze it) controlled by the dim parameter in torch.max.