import torch


def cat_arange(counts, dtype=torch.int32):
    """
    Concatenate results of multiple arange calls
    E.g.: cat_arange([2,1,3]) = [0, 1, 0, 0, 1, 2]
    Credits: https://stackoverflow.com/a/20033438

    :param torch.tensor counts: a 1D tensor
    :return: equivalent to torch.cat([torch.arange(c) for c in counts])
    """
    counts1 = counts[:-1].type(dtype)
    reset_index = torch.cumsum(counts1, dim=0).type(torch.int64)

    incr = torch.ones(counts.sum(), dtype=dtype, device=counts.device)
    incr[0] = 0
    incr[reset_index] = 1 - counts1

    # Reuse the incr array for the final result.
    return torch.cumsum(incr, dim=0)