add topk / maxk

[CarloLucibello]

We should add what is called elsewhere topk or maxk, e.g. https://pytorch.org/docs/stable/generated/torch.topk.html

In Base we have partialsort and partialsortperm doing something similar but limited to 1d arrays

[DhairyaLGandhi]

Those don't strictly appear to be NNlib's area to address though. We have topk in Metalhead.jl iirc

[DhairyaLGandhi]

We could also steal them from ResNetImageNet.jl

[CarloLucibello]

It's generally useful. Pytorch and tensorflow have it, deep graph library and pytorch geometric use it internally and I need it for https://github.com/CarloLucibello/GraphNeuralNetworks.jl. I'll look into and eventually port here those other implementations

[CarloLucibello]

I can't find topk in Metalhead.jl, nor in ResNetImageNet.jl

[DhairyaLGandhi]

https://github.com/FluxML/Metalhead.jl/blob/23df74c7df1bfca2ee66392a29bc3529deab9dbf/src/utils.jl#L73 (these functions should be ported back to master)

https://github.com/DhairyaLGandhi/ResNetImageNet.jl/blob/03497c26879d7462c1bb3c6199f54738c758fb84/src/utils.jl#L20

I agree that these are generally useful, but not a "neural network primitive". I can see it in a higher abstracted library like Metalhead, maybe Flux.

[darsnack]

I don't think they belong in Metalhead.jl as they are not specific to vision models. The advantage of NNlib is that packages unrelated to Flux can use it. topk is in the same category as losses for me, so Flux seems like the right place.

[CarloLucibello]

I consider it a differentiable transformation, a generalization of maximum, and prefer having it here for wider availability. We also have softmax here, which is a bit related as well