import torch


def repeat_arange(counts, dtype=torch.int32):
    """
    Repeat each element of arange multiple times
    E.g.: repeat_arange([2,1,3]) = [0, 0, 1, 2, 2, 2]

    :param counts: a 1D tensor having the same length of 'tensor'
    :return: equivalent to torch.cat([torch.tensor([v]).expand(n) for v, n in enumerate(counts)])
    """

    incr = torch.zeros(counts.sum(), dtype=dtype, device=counts.device)
    set_index = torch.cumsum(counts[:-1], dim=0).type(torch.int64)
    incr[set_index] = 1

    return torch.cumsum(incr, dim=0)