code: def _topk_on(
      graph, typestr, feat, k, descending, sortby, ntype_or_etype
  ):
      """Internal function to take graph-wise top-k node/edge features of
      field :attr:`feat` in :attr:`graph` ranked by keys at given
      index :attr:`sortby`. If :attr:`descending` is set to False, return the
      k smallest elements instead.
  
      Parameters
      ---------
      graph : DGLGraph
          The graph
      typestr : str
          'nodes' or 'edges'
      feat : str
          The feature field name.
      k : int
          The :math:`k` in "top-:math`k`".
      descending : bool
          Controls whether to return the largest or smallest elements,
           defaults to True.
      sortby : int
          The key index we sort :attr:`feat` on, if set to None, we sort
          the whole :attr:`feat`.
      ntype_or_etype : str, tuple of str
          Node/edge type.
  
      Returns
      -------
      sorted_feat : Tensor
          A tensor with shape :math:`(B, K, D)`, where
          :math:`B` is the batch size of the input graph.
      sorted_idx : Tensor
          A tensor with shape :math:`(B, K)`(:math:`(B, K, D)` if sortby
          is set to None), where
          :math:`B` is the batch size of the input graph, :math:`D`
          is the feature size.
  
  
      Notes
      -----
      If an example has :math:`n` nodes/edges and :math:`n<k`, in the first
      returned tensor the :math:`n+1` to :math:`k`th rows would be padded
      with all zero; in the second returned tensor, the behavior of :math:`n+1`
      to :math:`k`th elements is not defined.
      """
      _, batch_num_objs_attr, _ = READOUT_ON_ATTRS[typestr]
      data = getattr(graph, typestr)[ntype_or_etype].data
      if F.ndim(data[feat]) > 2:
          raise DGLError(
              "Only support {} feature `{}` with dimension less than or"
              " equal to 2".format(typestr, feat)
          )
      feat = data[feat]
      hidden_size = F.shape(feat)[-1]
      batch_num_objs = getattr(graph, batch_num_objs_attr)(
          ntype_or_etype
      )
      batch_size = len(batch_num_objs)
      length = max(max(F.asnumpy(batch_num_objs)), k)
      fill_val = -float("inf") if descending else float("inf")
      feat_ = F.pad_packed_tensor(
          feat, batch_num_objs, fill_val, l_min=k
      )  # (batch_size, l, d)
  
      if F.backend_name == "pytorch" and sortby is not None:
          # PyTorch's implementation of top-K
          keys = feat_[..., sortby]  # (batch_size, l)
          return _topk_torch(keys, k, descending, feat_)
      else:
          # Fallback to framework-agnostic implementation of top-K
          if sortby is not None:
              keys = F.squeeze(
                  F.slice_axis(feat_, -1, sortby, sortby + 1), -1
              )
              order = F.argsort(keys, -1, descending=descending)
          else:
              order = F.argsort(feat_, 1, descending=descending)
          topk_indices = F.slice_axis(order, 1, 0, k)
  
          if sortby is not None:
              feat_ = F.reshape(feat_, (batch_size * length, -1))
              shift = F.repeat(F.arange(0, batch_size) * length, k, -1)
              shift = F.copy_to(shift, F.context(feat))
              topk_indices_ = F.reshape(topk_indices, (-1,)) + shift
          else:
              feat_ = F.reshape(feat_, (-1,))
              shift = F.repeat(
                  F.arange(0, batch_size), k * hidden_size, -1
              ) * length * hidden_size + F.cat(
                  [F.arange(0, hidden_size)] * batch_size * k, -1
              )
              shift = F.copy_to(shift, F.context(feat))
              topk_indices_ = (
                  F.reshape(topk_indices, (-1,)) * hidden_size + shift
              )
          out = F.reshape(
              F.gather_row(feat_, topk_indices_), (batch_size, k, -1)
          )
          out = F.replace_inf_with_zero(out)
          return out, topk_indices
  
file_path: dgl/python/dgl/readout.py
imports: from __future__ import absolute_import
  from .None import backend
  from .base import dgl_warning, DGLError
  from .ops import segment
label_desc: Use image node_id 14 for calling a global function with example usage: _topk_on(graph, typestr, feat, k, descending, sortby, ntype_or_etype) and returns: _topk_torch, out,  topk_indices