import torch


def _compute_stacked_offsets(sizes, repeats):
    """Computes offsets to add to indices of stacked tensors (Tensorflow).

      When a set of tensors are stacked, the indices of those from the second on
      must be offset in order to be able to index into the stacked tensor. This
      computes those offsets.

      Args:
        sizes: A 1D `Tensor` of the sizes per graph.
        repeats: A 1D `Tensor` of the number of repeats per graph.

      Returns:
        A 1D `Tensor` containing the index offset per graph.
      """
    sizes = torch.tensor(sizes[:-1]).type(torch.long)
    offset_values = torch.cumsum(torch.cat([torch.zeros(1).long(), sizes], 0), dim=0)
    return torch.repeat_interleave(offset_values, repeats.long())