class_docstr: Find the k nearest neighbors
class_docstr_tok: ["Find","the","k","nearest","neighbors"]
docstr: Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch
docstr_tok: ["Adapted","from","https",":","\/\/github.com\/yanx27\/Pointnet_Pointnet2_pytorch"]
code: def forward(self, pos, centroids):
      """
      Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch
      """
      center_pos = index_points(pos, centroids)
      sqrdists = square_distance(center_pos, pos)
      group_idx = sqrdists.argsort(dim=-1)[:, :, : self.n_neighbor]
      return group_idx
  
imports: import dgl
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl.geometry import farthest_point_sampler
label_desc: Use image node_id 2 for calling the KNearNeighbors obj's underlying member method code with example usage: obj.forward(pos, centroids) and returns: group_idx