def _stack_concatenated_datasets(concat_dataset):
    """Extract and stack indices of different data sets from `concat_dataset`.

    Each data set is represented by a complete range of indices
    starting from 0...len(data set) for the first data set.
    The ranges of the following data sets build on
    the range of the corresponding previous data set so that
    the indices of the lists are cumulative::

        datasets = [[0 ... len(data_1) - 1],
                    [len(data_1) ... len(data_1) - 1], ... ]

    Args:
        concat_dataset (ConcatDataset): data set to sample from

    Returns:
        list: list of lists of data set indices

    Example:
        >>> import torch
        >>> from torch.utils.data import (TensorDataset, ConcatDataset)
        >>> data_1 = TensorDataset(torch.Tensor([0, 1, 2, 3]))
        >>> data_2 = TensorDataset(torch.Tensor([0, 1, 2]))
        >>> concat_dataset = ConcatDataset((data_1, data_2))
        >>> concat_dataset.cumulative_sizes
        [4, 7]
        >>> _stack_concatenated_datasets(concat_dataset)
        [[0, 1, 2, 3], [4, 5, 6]]

    """
    datasets = []
    start_idx = 0
    for upper_edge in concat_dataset.cumulative_sizes:
        datasets += [list(range(start_idx, upper_edge))]
        start_idx = upper_edge
    return datasets
