code: def load_split_from_tfds_builder(
      builder,
      batch_size,
      split,
      preprocess_example=None,
      augment_train_example=None,
      postprocess_batch=None,
      shuffle_buffer_size=None,
      shuffle_seed=0,
      cache=True,
  ):
      """Loads a split from a dataset using TensorFlow Datasets compatible builder.
  
      Args:
        builder: tfds.core.DatasetBuilder; A TFDS compatible dataset builder.
        batch_size: int; The batch size returned by the data pipeline.
        split: str; Name of  the split to be loaded.
        preprocess_example: function; A function that given an example, returns the
          preprocessed example. Note that the preprocessing is done BEFORE caching
          to re-use them.
        augment_train_example: A function that given a train example returns the
          augmented example. Note that this function is applied AFTER caching and
          repeat to get true randomness.
        postprocess_batch: function; A function that given a batch, returns the
          postprocessed batch.
        shuffle_buffer_size: int; Size of the tf.data.dataset shuffle buffer.
        shuffle_seed: int; Seed for shuffling the training data.
        cache: bool; Whether to cache dataset in memory.
  
      Returns:
        A `tf.data.Dataset`, and dataset information.
      """
      # Prepare map functions.
      preprocess_example = preprocess_example or (lambda ex: ex)
      augment_train_example = augment_train_example or (lambda ex: ex)
      postprocess_batch = postprocess_batch or (lambda ex: ex)
      shuffle_buffer_size = shuffle_buffer_size or (8 * batch_size)
  
      # Download dataset:
      builder.download_and_prepare()
  
      # Each host is responsible for a fixed subset of data.
      data_range = tfds.even_splits(split, jax.process_count())[
          jax.process_index()
      ]
      ds = builder.as_dataset(split=data_range, shuffle_files=False)
      options = tf.data.Options()
      options.threading.private_threadpool_size = 48
      ds = ds.with_options(options)
  
      # Applying preprocessing before `ds.cache()` to re-use it.
      ds = ds.map(
          preprocess_example,
          num_parallel_calls=tf.data.experimental.AUTOTUNE,
      )
      # Caching.
      if cache:
          ds = ds.cache()
  
      if "train" in split:
          # First repeat then batch.
          ds = ds.repeat()
          # Augmentation should be done after repeat for true randomness.
          ds = ds.map(
              augment_train_example,
              num_parallel_calls=tf.data.experimental.AUTOTUNE,
          )
          # Shuffle after augmentation to avoid loading uncropped images into buffer:
          ds = ds.shuffle(shuffle_buffer_size, seed=shuffle_seed)
          ds = ds.batch(batch_size, drop_remainder=True)
  
      else:
          # First batch then repeat.
          ds = ds.batch(batch_size, drop_remainder=False)
          ds = ds.repeat()
  
      ds = ds.map(
          postprocess_batch,
          num_parallel_calls=tf.data.experimental.AUTOTUNE,
      )
      ds = ds.prefetch(tf.data.experimental.AUTOTUNE)
      return ds, builder.info
  
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 10 for calling a global function with example usage: load_split_from_tfds_builder(builder, batch_size, split, preprocess_example, augment_train_example, postprocess_batch, shuffle_buffer_size, shuffle_seed, cache) and returns: ds,  builder