code: def iterator_from_dataset(
      dataset: tf.data.Dataset,
      batch_size: int,
      repeat: bool = True,
      prefetch_size: int = 0,
      devices: Optional[Sequence[Any]] = None,
  ):
      """Create a data iterator that returns JAX arrays from a TF dataset.
  
      Args:
        dataset: the dataset to iterate over.
        batch_size: the batch sizes the iterator should return.
        repeat: whether the iterator should repeat the dataset.
        prefetch_size: the number of batches to prefetch to device.
        devices: the devices to prefetch to.
  
      Returns:
        An iterator that returns data batches.
      """
      if repeat:
          dataset = dataset.repeat()
  
      if batch_size > 0:
          dataset = dataset.batch(batch_size)
          it = map(prepare_tf_data, dataset)
      else:
          it = map(prepare_tf_data_unbatched, dataset)
  
      if prefetch_size > 0:
          it = jax_utils.prefetch_to_device(it, prefetch_size, devices)
  
      return it
  
imports: import abc
  import functools
  import itertools
  from typing import Any, Iterable, Optional, Sequence, Union
  from absl import logging
  from flax import jax_utils
  import jax
  import numpy
  import tensorflow
  from nerfies import camera
  from nerfies import gpath
  from nerfies import image_utils
  from nerfies import tf_camera
  from nerfies import utils
  from tensorflow.python.data.util import nest
label_desc: Use image node_id 5 for calling a global function with example usage: iterator_from_dataset(dataset, batch_size, repeat, prefetch_size, devices) and returns: it