code: def get_dists(g, nbrs, use_gt):
      k = nbrs.shape[1]
      src_id = nbrs[:, 1:].reshape(-1)
      dst_id = nbrs[:, 0].repeat(k - 1)
      eids = g.edge_ids(src_id, dst_id)
      if use_gt:
          new_dists = (
              (1 - g.edata["labels_edge"][eids])
              .reshape(-1, k - 1)
              .float()
          )
      else:
          new_dists = g.edata["prob_conn"][eids, 0].reshape(-1, k - 1)
      ind = torch.argsort(new_dists, 1)
      offset = torch.LongTensor(
          (nbrs[:, 0] * (k - 1)).repeat(k - 1).reshape(-1, k - 1)
      ).to(g.device)
      ind = ind + offset
      nbrs = torch.LongTensor(nbrs).to(g.device)
      new_nbrs = torch.take(nbrs[:, 1:], ind)
      new_dists = torch.cat(
          [
              torch.zeros((new_dists.shape[0], 1)).to(g.device),
              new_dists,
          ],
          dim=1,
      )
      new_nbrs = torch.cat(
          [
              torch.arange(new_nbrs.shape[0]).view(-1, 1).to(g.device),
              new_nbrs,
          ],
          dim=1,
      )
      return (
          new_nbrs.cpu().detach().numpy(),
          new_dists.cpu().detach().numpy(),
      )
  
file_path: dgl/examples/pytorch/hilander/utils/deduce.py
imports: import dgl
  import numpy
  import torch
  from sklearn import mixture
  from .density import density_to_peaks, density_to_peaks_vectorize
label_desc: Use image node_id 5 for calling a global function with example usage: get_dists(g, nbrs, use_gt) and returns: new_nbrs,  new_dists