import torch as th


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)`.
    """
    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