class_docstr: Ball Query - Find the neighbors with-in a fixed radius
class_docstr_tok: ["Ball","Query","-","Find","the","neighbors","with-in","a","fixed","radius"]
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
      """
      device = pos.device
      B, N, _ = pos.shape
      center_pos = index_points(pos, centroids)
      _, S, _ = center_pos.shape
      group_idx = (
          torch.arange(N, dtype=torch.long)
          .to(device)
          .view(1, 1, N)
          .repeat([B, S, 1])
      )
      sqrdists = square_distance(center_pos, pos)
      group_idx[sqrdists > self.radius**2] = N
      group_idx = group_idx.sort(dim=-1)[0][:, :, : self.n_neighbor]
      group_first = (
          group_idx[:, :, 0]
          .view(B, S, 1)
          .repeat([1, 1, self.n_neighbor])
      )
      mask = group_idx == N
      group_idx[mask] = group_first[mask]
      return group_idx
  
imports: import dgl
  import dgl.function
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl.geometry import farthest_point_sampler
  from torch.autograd import Variable
label_desc: Use image node_id 2 for calling the FixedRadiusNearNeighbors obj's underlying member method code with example usage: obj.forward(pos, centroids) and returns: group_idx