code: def create_split(
      dataset_builder,
      batch_size,
      train,
      dtype=tf.float32,
      image_size=IMAGE_SIZE,
      cache=False,
      shuffle_buffer_size=2_000,
      prefetch=10,
  ):
      """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.
        image_size: The target size of the images.
        cache: Whether to cache the dataset.
        shuffle_buffer_size: Size of the shuffle buffer.
        prefetch: Number of items to prefetch in the dataset.
      Returns:
        A `tf.data.Dataset`.
      """
      if train:
          train_examples = dataset_builder.info.splits[
              "train"
          ].num_examples
          split_size = train_examples // jax.process_count()
          start = jax.process_index() * split_size
          split = f"train[{start}:{start + split_size}]"
      else:
          validate_examples = dataset_builder.info.splits[
              "validation"
          ].num_examples
          split_size = validate_examples // jax.process_count()
          start = jax.process_index() * split_size
          split = f"validation[{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(),
          },
      )
      options = tf.data.Options()
      options.experimental_threading.private_threadpool_size = 48
      ds = ds.with_options(options)
  
      if cache:
          ds = ds.cache()
  
      if train:
          ds = ds.repeat()
          ds = ds.shuffle(shuffle_buffer_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(prefetch)
  
      return ds
  
file_path: flax/examples/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, shuffle_buffer_size, prefetch) and returns: ds, dict