keras/src/trainers/data_adapters/generator_data_adapter.py/peek_and_restore
def peek_and_restore(generator):
    batches = list(
        itertools.islice(
            generator, data_adapter_utils.NUM_BATCHES_FOR_TENSOR_SPEC
        )
    )
    return batches, itertools.chain(batches, generator)