keras/src/trainers/data_adapters/torch_data_loader_adapter.py/TorchDataLoaderAdapter/get_jax_iterator
class TorchDataLoaderAdapter:    def get_jax_iterator(self):
        # We use numpy as an intermediary because the conversion
        # torch -> numpy -> jax is faster than torch -> jax.
        return data_adapter_utils.get_jax_iterator(self.get_numpy_iterator())