def _split_user_def_task_label(
    datasets, task_labels: Optional[Union[int, Sequence[int], Sequence[Sequence[int]]]]
) -> List[Optional[Union[int, Sequence[int]]]]:
    """
    Given a datasets list and the user-defined list of task labels,
    returns the task labels list of each dataset.

    This internal utility is mainly used to manage the different ways
    in which the user can define the task labels:
    - As a single task label for all exemplars of all datasets
    - A single list of length equal to the sum of the lengths of all datasets
    - A list containing, for each dataset, one element between:
        - a list, defining the task labels of each exemplar of a that dataset
        - an int, defining the task label of all exemplars of a that dataset

    :param datasets: The list of datasets.
    :param task_labels: The user-defined task labels. Can be None, in which
        case a list of None will be returned.
    :returns: A list containing as many elements as the input `datasets`.
        Each element is either a list of task labels or None. If None
        (because `task_labels` is None), this means that the task labels
        should be retrieved by traversing each dataset.
    """
    t_labels = []
    idx_start = 0
    for dd_idx, dd in enumerate(datasets):
        end_idx = idx_start + len(dd)
        dataset_t_label: Optional[Union[int, Sequence[int]]]
        if task_labels is None:
            # No task label set
            dataset_t_label = None
        elif isinstance(task_labels, int):
            # Single integer (same label for all instances)
            dataset_t_label = task_labels
        elif isinstance(task_labels[0], int):
            # Single task labels sequence
            # (to be split across concatenated datasets)
            dataset_t_label = task_labels[idx_start:end_idx]  # type: ignore
        elif len(task_labels[dd_idx]) == len(dd):  # type: ignore
            # One sequence per dataset
            dataset_t_label = task_labels[dd_idx]
        else:
            raise ValueError("The task_labels parameter has an invalid format.")
        t_labels.append(dataset_t_label)

        idx_start = end_idx
    return t_labels

    Given a datasets list and the user-defined list of task labels,
    returns the task labels list of each dataset.

    This internal utility is mainly used to manage the different ways
    in which the user can define the task labels:
    - As a single task label for all exemplars of all datasets
    - A single list of length equal to the sum of the lengths of all datasets
    - A list containing, for each dataset, one element between:
        - a list, defining the task labels of each exemplar of a that dataset
        - an int, defining the task label of all exemplars of a that dataset

    :param datasets: The list of datasets.
    :param task_labels: The user-defined task labels. Can be None, in which
        case a list of None will be returned.
    :returns: A list containing as many elements as the input `datasets`.
        Each element is either a list of task labels or None. If None
        (because `task_labels` is None), this means that the task labels
        should be retrieved by traversing each dataset.
    
avalanche/benchmarks/utils/utils.py