import itertools


def _cells_initialize_states(cells, batch_axis, **kwargs):
    """
    Given a list of cells, initialize the states of the individual
    cells together by doing the equivalent of
    sum([c.initialize_states(**kwargs) for c in cells], []),
    but using itertools instead because it's much faster

    Arguments:
    ----------
    cells: list containing cell objects
    """

    return list(itertools.chain(
        *[c.initialize_states(batch_axis, **kwargs) for c in cells]))