code: def process_record_dataset(
      dataset,
      is_training,
      batch_size,
      num_epochs,
      shuffle_buffer=None,
      dtype=tf.float32,
      datasets_num_private_threads=None,
      drop_remainder=False,
      tf_data_experimental_slack=False,
      prefetch_batchs=tf.data.experimental.AUTOTUNE,
  ):
      """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.
        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.
        prefetch_batchs: The number of batchs to prefetch.
  
      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.
          if shuffle_buffer is not None:
              dataset = dataset.shuffle(buffer_size=shuffle_buffer)
          # Repeats the dataset for the number of epochs to train.
          dataset = dataset.repeat(num_epochs)
  
      one_hot = False
      # TODO(swang): Support one-hot encoding?
      # num_classes = FLAGS.num_classes
      # if FLAGS.label_smoothing and FLAGS.label_smoothing > 0:
      #  one_hot = True
  
      num_classes = NUM_CLASSES
  
      map_fn = functools.partial(
          preprocess_parsed_example,
          is_training=is_training,
          dtype=dtype,
          num_classes=num_classes,
          one_hot=one_hot,
      )
  
      # Parses the raw records into images and labels.
      dataset = dataset.map(
          map_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE
      )
      dataset = dataset.batch(batch_size, drop_remainder=drop_remainder)
  
      # Operations between the final prefetch and the get_next call to the iterator
      # will happen synchronously during run time. We prefetch here again to
      # background all of the above processing work and keep it out of the
      # critical training path. Setting buffer_size to tf.data.experimental.AUTOTUNE
      # allows DistributionStrategies to adjust how many batches to fetch based
      # on how many devices are present.
      dataset = dataset.prefetch(buffer_size=prefetch_batchs)
  
      options = tf.data.Options()
      options.experimental_slack = tf_data_experimental_slack
      dataset = dataset.with_options(options)
  
      return dataset
  
file_path: ray/release/air_tests/air_benchmarks/mlperf-train/tf_utils.py
imports: import tensorflow
  import functools
  import logging
label_desc: Use image node_id 1 for calling a global function with example usage: process_record_dataset(dataset, is_training, batch_size, num_epochs, shuffle_buffer, dtype, datasets_num_private_threads, drop_remainder, tf_data_experimental_slack, prefetch_batchs) and returns: dataset