code: def _topk_torch(keys, k, descending, x):
      """Internal function to take graph-wise top-k node/edge features according to
      the rank given by keys, this function is PyTorch only.
  
      Parameters
      ----------
      keys : Tensor
          The key for ranking.
      k : int
          The :math:`k` in "top-:math:`k`".
      descending : bool
          Indicates whether to return the feature corresponding to largest or
          smallest elements.
      x : Tensor
          The padded feature with shape (batch, max_len, *)
  
      Returns
      -------
      sorted_feat : Tensor
          A tensor with shape :math:`(batch, k, *)`.
      sorted_idx : Tensor
          A tensor with shape :math:`(batch, k)`.
      """
      import torch as th
  
      batch_size, max_len = x.shape[0], x.shape[1]
      topk_indices = keys.topk(k, -1, largest=descending)[
          1
      ]  # (batch_size, k)
      x = x.view((batch_size * max_len), -1)
      shift = (
          th.arange(0, batch_size, device=x.device).view(batch_size, 1)
          * max_len
      )
      topk_indices_ = topk_indices + shift
      x = x[topk_indices_].view(batch_size, k, -1)
      return th.masked_fill(x, th.isinf(x), 0), 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 13 for calling a global function with example usage: _topk_torch(keys, k, descending, x) and returns: th,  topk_indices