import torch


def offsets_from_counts(counts):
    """Creates a tensor of offsets from the given tensor of counts

    Parameters
    ----------
    counts : torch.Tensor
        1-d tensor representing the counts in a ragged array

    Returns
    -------
    torch.Tensor
        A 1-d tensor representing the offsets in a ragged array.
        Its length is one plus the length of the input `counts` array.
    """
    if isinstance(counts, torch.Tensor):
        device = counts.device
    else:
        device = None
    counts = torch.as_tensor(counts, dtype=torch.int64, device=device)
    offsets = counts.new_zeros(counts.shape[0] + 1)
    torch.cumsum(counts, 0, out=offsets[1:])
    return offsets