def topk_mask(xs: TensorType["Batch", "Vocab"], k: int):
    """
    Takes batched distribution over tokens and masks out scores for tokens
    that are not in the top k for that distribution.
    """

    # Get topk per distribution
    # For each dist, getting last value gives k-th largest
    mintop = torch.topk(xs, k)[0][:, -1].unsqueeze(-1)
    return torch.where(xs < mintop, -np.inf * torch.ones_like(xs), xs)

    Takes batched distribution over tokens and masks out scores for tokens
    that are not in the top k for that distribution.
    
promptist/trlx/trlx/utils/__init__.py