code: def process_record_dataset(
      dataset,
      is_training,
      batch_size,
      shuffle_buffer,
      parse_record_fn,
      num_epochs=1,
      dtype=tf.float32,
      datasets_num_private_threads=None,
      drop_remainder=False,
      tf_data_experimental_slack=False,
  ):
      """Given a Dataset with raw records, return an iterator over the records.
  
      Args:
        dataset: A Dataset representing raw records
        is_training: A boolean denoting whether the input is for training.
        batch_size: The number of samples per batch.
        shuffle_buffer: The buffer size to use when shuffling records. A larger
          value results in better randomness, but smaller values reduce startup
          time and use less memory.
        parse_record_fn: A function that takes a raw record and returns the
          corresponding (image, label) pair.
        num_epochs: The number of epochs to repeat the dataset.
        dtype: Data type to use for images/features.
        datasets_num_private_threads: Number of threads for a private
          threadpool created for all datasets computation.
        drop_remainder: A boolean indicates whether to drop the remainder of the
          batches. If True, the batch dimension will be static.
        tf_data_experimental_slack: Whether to enable tf.data's
          `experimental_slack` option.
  
      Returns:
        Dataset of (image, label) pairs ready for iteration.
      """
      # Defines a specific size thread pool for tf.data operations.
      if datasets_num_private_threads:
          options = tf.data.Options()
          options.experimental_threading.private_threadpool_size = (
              datasets_num_private_threads
          )
          dataset = dataset.with_options(options)
          logging.info(
              "datasets_num_private_threads: %s",
              datasets_num_private_threads,
          )
  
      if is_training:
          # Shuffles records before repeating to respect epoch boundaries.
          dataset = dataset.shuffle(buffer_size=shuffle_buffer)
          # Repeats the dataset for the number of epochs to train.
          dataset = dataset.repeat()
  
      # Parses the raw records into images and labels.
      dataset = dataset.map(
          lambda value: parse_record_fn(value, is_training, dtype),
          num_parallel_calls=tf.data.experimental.AUTOTUNE,
      )
      dataset = dataset.batch(batch_size, drop_remainder=drop_remainder)
  
      dataset = dataset.prefetch(
          buffer_size=tf.data.experimental.AUTOTUNE
      )
  
      options = tf.data.Options()
      options.experimental_slack = tf_data_experimental_slack
      dataset = dataset.with_options(options)
  
      return dataset
  
file_path: autodist/examples/benchmark/utils/imagenet_preprocessing.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  from absl import logging
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: process_record_dataset(dataset, is_training, batch_size, shuffle_buffer, parse_record_fn, num_epochs, dtype, datasets_num_private_threads, drop_remainder, tf_data_experimental_slack) and returns: dataset