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