import torch


def batch_randperm(n_batch: int, n: int, device: str = 'cpu') -> torch.Tensor:
    """
    Returns multiple random permutations.
    :param n_batch: Number of permutations.
    :param n: Length of permutations.
    :param device: PyTorch Device to put the permutations on.
    :return: Returns a torch.Tensor of integer type of shape [n_batch, n] containing the n_batch permutations.
    """
    # batched randperm:
    # https://discuss.pytorch.org/t/batched-shuffling-of-feature-vectors/30188/4
    # https://github.com/pytorch/pytorch/issues/42502
    return torch.stack([torch.randperm(n, device=device) for i in range(n_batch)], dim=0)