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
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.
    
python/dgl/readout.py