docstr: Shard and prefetch batches on device.
  
  Most of the implementation has been borrowed from
  `flax.jax_utils.prefetch_to_device`
  
  This utility takes an iterator and returns a new iterator which fills an
  on device prefetch buffer. Eager prefetching can improve the performance
  of training loops significantly by overlapping compute and data
  transfer.
docstr_tok: ["Shard","and","prefetch","batches","on","device",".","Most","of","the","implementation","has","been","borrowed","from","`","flax.jax_utils.prefetch_to_device","`","This","utility","takes","an","iterator","and","returns","a","new","iterator","which","fills","an","on","device","prefetch","buffer",".","Eager","prefetching","can","improve","the","performance","of","training","loops","significantly","by","overlapping","compute","and","data","transfer","."]
code: def _prefetch_numpy_iterator(self, numpy_iterator):
      """Shard and prefetch batches on device.
  
      Most of the implementation has been borrowed from
      `flax.jax_utils.prefetch_to_device`
  
      This utility takes an iterator and returns a new iterator which fills an
      on device prefetch buffer. Eager prefetching can improve the performance
      of training loops significantly by overlapping compute and data
      transfer.
      """
      queue = collections.deque()
  
      # If you're training on GPUs, 2 is generally the best choice because
      # this guarantees that you can overlap a training step on GPU with a
      # data prefetch step on CPU.
      def enqueue(n=2):
          for data in itertools.islice(numpy_iterator, n):
              queue.append(_distribute_data(data))
  
      enqueue(n=2)  # TODO: should we make `n` configurable?
      while queue:
          yield queue.popleft()
          enqueue(1)
  
imports: import collections
  import itertools
  from functools import partial
  import jax
  import numpy
  import tree
  from keras import backend
  from keras import callbacks
  from keras import ops
  from keras import optimizers
  from keras.backend import distribution_lib
  from keras.distribution import distribution_lib
  from keras.trainers import trainer
  from keras.trainers.data_adapters import data_adapter_utils
  from keras.trainers.epoch_iterator import EpochIterator
  from keras.utils import traceback_utils
label_desc: Use image node_id 2 for calling the JAXEpochIterator obj's underlying member method code with example usage: obj._prefetch_numpy_iterator(numpy_iterator) without return types