code: def _neighbor_sort(
      scores: Tensor,
      end_n_ids: Tensor,
      in_degrees: Tensor,
      cum_in_degrees: Tensor,
  ):
      """Sort edge scores for each node"""
      num_nodes, max_in_degree = in_degrees.size(0), int(
          in_degrees.max().item()
      )
  
      # Compute the index for dense score matrix with size (N x D_{max})
      # Note that the end_n_ids here is the end_node tensor in dgl graph,
      # which is not grouped by its node id (i.e. in this form: 0,0,1,1,1,...,N,N).
      # Thus here we first sort the end_node tensor to make it easier to compute
      # indexs in dense edge score matrix. Since we will need the original order
      # for following gspmm and gsddmm operations, we also keep the reverse mapping
      # (the reverse_perm) here.
      end_n_ids, perm = torch.sort(end_n_ids)
      scores = scores[perm]
      _, reverse_perm = torch.sort(perm)
  
      index = torch.arange(
          end_n_ids.size(0), dtype=torch.long, device=scores.device
      )
      index = (index - cum_in_degrees[end_n_ids]) + (
          end_n_ids * max_in_degree
      )
      index = index.long()
  
      dense_scores = scores.new_full(
          (num_nodes * max_in_degree,), torch.finfo(scores.dtype).min
      )
      dense_scores[index] = scores
      dense_scores = dense_scores.view(num_nodes, max_in_degree)
  
      sorted_dense_scores, dense_reverse_perm = dense_scores.sort(
          dim=-1, descending=True
      )
      _, dense_reverse_perm = torch.sort(dense_reverse_perm, dim=-1)
      dense_reverse_perm = dense_reverse_perm + cum_in_degrees.view(
          -1, 1
      )
      dense_reverse_perm = dense_reverse_perm.view(-1)
      cumsum_sorted_dense_scores = sorted_dense_scores.cumsum(
          dim=-1
      ).view(-1)
      sorted_dense_scores = sorted_dense_scores.view(-1)
      arange_vec = torch.arange(
          1,
          max_in_degree + 1,
          dtype=torch.long,
          device=end_n_ids.device,
      )
      arange_vec = torch.repeat_interleave(
          arange_vec.view(1, -1), num_nodes, dim=0
      ).view(-1)
  
      valid_mask = sorted_dense_scores != torch.finfo(scores.dtype).min
      sorted_scores = sorted_dense_scores[valid_mask]
      cumsum_sorted_scores = cumsum_sorted_dense_scores[valid_mask]
      arange_vec = arange_vec[valid_mask]
      dense_reverse_perm = dense_reverse_perm[valid_mask].long()
  
      return (
          sorted_scores,
          cumsum_sorted_scores,
          arange_vec,
          reverse_perm,
          dense_reverse_perm,
      )
  
file_path: dgl/examples/pytorch/hgp_sl/functions.py
imports: import dgl
  import torch
  from dgl.backend import astype
  from dgl.base import ALL, is_all
  from dgl.heterograph_index import HeteroGraphIndex
  from dgl.sparse import _gsddmm, _gspmm
  from torch import Tensor
  from torch.autograd import Function
label_desc: Use image node_id 1 for calling a global function with example usage: _neighbor_sort(scores, end_n_ids, in_degrees, cum_in_degrees) and returns: sorted_scores,  cumsum_sorted_scores,  arange_vec,  reverse_perm,  dense_reverse_perm