code: def tune_dataset(
      dataset,
      batch_size=None,
      shuffle_size=None,
      preprocess_fn=None,
      repeat_count=-1,
  ):
      """Sets generally recommended parameters for a `tf.data.Dataset`.
  
      Args:
        dataset: `tf.data.Dataset`-like instance to be tuned according to this
          functions arguments.
        batch_size: Python `int` representing the number of elements in each
          minibatch.
        shuffle_size: Python `int` representing the number of elements to shuffle
          (at a time).
        preprocess_fn: Python `callable` applied to each item in `dataset`.
        repeat_count: Python `int`, representing the number of times the dataset
          should be repeated. The default behavior (`repeat_count = -1`) is for the
          dataset to be repeated indefinitely. If `repeat_count is None` repeat is
          "off;" note that this is a deviation from `tf.data.Dataset.repeat` which
          interprets `None` as "repeat indefinitely".
          Default value: `-1` (i.e., repeat indefinitely).
  
      Returns:
        tuned_dataset: `tf.data.Dataset` instance tuned according to this functions
          arguments.
  
      #### Example
  
      ```python
      [train_dataset, eval_dataset], datasets_info = tfds.load(
           name='mnist',
           split=['train', 'test'],
           with_info=True,
           as_supervised=True,
           shuffle_files=True)
  
      def _preprocess(image, label):
        image = tf.cast(image, dtype=tf.int32)
        u = tf.random.uniform(shape=tf.shape(image), maxval=256, dtype=image.dtype)
        image = tf.cast(u < image, dtype=tf.float32)   # Randomly binarize.
        return image, label
  
      # TODO(b/144500779): Cant use `jit_compile=True`.
      @tf.function(autograph=False)
      def one_step(iter):
        x, y = next(iter)
        return tf.reduce_mean(x)
  
      ds = tune_dataset(
          train_dataset,
          batch_size=32,
          shuffle_size=int(datasets_info.splits['train'].num_examples / 7),
          preprocess_fn=_preprocess)
      it = iter(ds)
      [one_step(it)]*3  # Build graph / burn-in.
      %time one_step(it)
      ```
  
      """
      # https://www.tensorflow.org/guide/data_performance
      # The order of these builder arguments matters.
      # The following has been tuned using the snippet above.
      if preprocess_fn is not None:
          dataset = dataset.map(
              preprocess_fn,
              num_parallel_calls=tf.data.experimental.AUTOTUNE,
          )
      dataset = dataset.cache()
      if shuffle_size is not None:
          dataset = dataset.shuffle(shuffle_size)
      # Repeat must follow shuffle, else samples could be dropped if we only
      # complete one epoch. E.g.,
      # ds = tf.data.Dataset.range(10).repeat(2).shuffle(5)
      # seen = set(e.numpy() for i, e in enumerate(ds) if i < 10)  # One epoch.
      # assert len(seen) == 10
      if repeat_count is not None:
          dataset = dataset.repeat(repeat_count)
      if batch_size is not None:
          dataset = dataset.batch(batch_size, drop_remainder=True)
      dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
      return dataset
  
imports: import collections
  import contextlib
  import functools
  import sys
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 2 for calling a global function with example usage: tune_dataset(dataset, batch_size, shuffle_size, preprocess_fn, repeat_count) and returns: dataset