keras/src/backend/jax/trainer.py/JAXEpochIterator/_get_iterator
class JAXEpochIterator:    def _get_iterator(self):
        return self._prefetch_numpy_iterator(
            self.data_adapter.get_jax_iterator()
        )