code: def create_split(
      dataset_builder: tfds.core.DatasetBuilder,
      batch_size: int,
      train: bool,
      dtype: tf.DType = tf.float32,
      image_size: int = IMAGE_SIZE,
      cache: bool = False,
  ):
      """Creates a split from the ImageNet dataset using TensorFlow Datasets.
  
      Args:
        dataset_builder: TFDS dataset builder for ImageNet.
        batch_size: the batch size returned by the data pipeline.
        train: Whether to load the train or evaluation split.
        dtype: data type of the image (default: float32).
        image_size: The target size of the images (default: 224).
        cache: Whether to cache the dataset (default: False).
      Returns:
        A `tf.data.Dataset`.
      """
      if train:
          train_size = dataset_builder.info.splits["train"].num_examples
          split_size = train_size // jax.host_count()
          start = jax.host_id() * split_size
          split = "train[{}:{}]".format(start, start + split_size)
      else:
          validation_size = dataset_builder.info.splits[
              "validation"
          ].num_examples
          split_size = validation_size // jax.host_count()
          start = jax.host_id() * split_size
          split = "validation[{}:{}]".format(start, start + split_size)
  
      def _decode_example(example):
          if train:
              image = preprocess_for_train(
                  example["image"], dtype, image_size
              )
          else:
              image = preprocess_for_eval(
                  example["image"], dtype, image_size
              )
          return {"image": image, "label": example["label"]}
  
      ds = dataset_builder.as_dataset(
          split=split, decoders={"image": tfds.decode.SkipDecoding()}
      )
      ds.options().experimental_threading.private_threadpool_size = 48
      ds.options().experimental_threading.max_intra_op_parallelism = 1
  
      if cache:
          ds = ds.cache()
  
      if train:
          ds = ds.repeat()
          ds = ds.shuffle(16 * batch_size, seed=0)
  
      ds = ds.map(
          _decode_example,
          num_parallel_calls=tf.data.experimental.AUTOTUNE,
      )
  
      ds = ds.batch(batch_size, drop_remainder=True)
  
      if not train:
          ds = ds.repeat()
  
      ds = ds.prefetch(10)
  
      return ds
  
file_path: elegy/examples/need-fixing/imagenet/input_pipeline.py
imports: import jax
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 9 for calling a global function with example usage: create_split(dataset_builder, batch_size, train, dtype, image_size, cache) and returns: ds, dict