code: def create_input_iter(
      dataset_builder,
      batch_size,
      image_size,
      dtype,
      train,
      cache,
      shuffle_buffer_size,
      prefetch,
  ):
      ds = input_pipeline.create_split(
          dataset_builder,
          batch_size,
          image_size=image_size,
          dtype=dtype,
          train=train,
          cache=cache,
          shuffle_buffer_size=shuffle_buffer_size,
          prefetch=prefetch,
      )
      it = map(prepare_tf_data, ds)
      it = jax_utils.prefetch_to_device(it, 2)
      return it
  
imports: import functools
  import time
  from typing import Any
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from flax import jax_utils
  from flax.training import checkpoints
  from flax.training import common_utils
  from flax.training import dynamic_scale
  from flax.training import train_state
  import jax
  from jax import lax
  import jax.numpy
  from jax import random
  import ml_collections
  import optax
  import tensorflow
  import tensorflow_datasets
  import input_pipeline
  import models
label_desc: Use image node_id 9 for calling a global function with example usage: create_input_iter(dataset_builder, batch_size, image_size, dtype, train, cache, shuffle_buffer_size, prefetch) and returns: it