keras/src/trainers/data_adapters/data_adapter_utils.py/get_jax_iterator
def get_jax_iterator(iterable):
    from keras.src.backend.jax.core import convert_to_tensor

    for batch in iterable:
        yield tree.map_structure(convert_to_tensor, batch)