def _restore_dataset_from_list(
    dataset_as_list, dataset_type_spec, original_dataset
):
    """Restore the dataset from the list of arrays."""
    if dataset_type_spec in [tuple, list, tf.data.Dataset] or is_torch_dataset(
        original_dataset
    ):
        # Save structure by taking the first element.
        element_spec = dataset_as_list[0]
        # Flatten each element.
        dataset_as_list = [tree.flatten(sample) for sample in dataset_as_list]
        # Combine respective elements at all indices.
        dataset_as_list = [np.array(sample) for sample in zip(*dataset_as_list)]
        # Recreate the original structure of elements.
        dataset_as_list = tree.pack_sequence_as(element_spec, dataset_as_list)
        # Turn lists to tuples as tf.data will fail on lists.
        return tree.traverse(
            lambda x: tuple(x) if isinstance(x, list) else x,
            dataset_as_list,
            top_down=False,
        )

    return dataset_as_list
Restore the dataset from the list of arrays.
keras/src/utils/dataset_utils.py