code: def prefetch_iterator(it, n):
      """Prefetches batches from an iterator.
  
      Runs iterator `it` ahead for `n` steps.
  
      Adapted from big_vision:
      https://github.com/google-research/big_vision/blob/main/big_vision/input_pipeline.py.
  
      Args:
        it: Iterator
        n: Number of steps to prefect for.
  
      Yields:
        Original items from the iterator which have been prefetched.
      """
      if not n:
          yield from it
          return
      queue = collections.deque()
  
      def enqueue(
          n_steps,
      ):  # Enqueues *up to* `n` elements from the iterator.
          for data in itertools.islice(it, n_steps):
              queue.append(data)
  
      enqueue(n)  # Fill up the buffer.
      while queue:
          yield queue.popleft()
          enqueue(1)
  
imports: import collections
  import dataclasses
  import functools
  import itertools
  from typing import Any, Callable, Dict, Iterator, Optional, Sequence, Union
  from absl import logging
  from flax.training import common_utils
  import jax
  import jax.numpy
  import numpy
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 4 for calling a global function with example usage: prefetch_iterator(it, n) without return types