This issue is a sibling of a similar issue for k-means: #26585 with similar purpose but likely different constraints.

In particular an efficient implementation of k-NN on the GPU would require:

• torch.cdist
• torch.topk being discussed at:
  • RFC: add topk and / or argpartition data-apis/array-api#629