code: def build_tf_dataset(
      filenames,
      batch_size,
      num_images_per_epoch,
      num_epochs,
      shuffle_buffer=None,
      is_training=True,
      dtype=tf.float32,
      datasets_num_private_threads=None,
      input_context=None,
      drop_remainder=False,
      tf_data_experimental_slack=False,
      dataset_cache=False,
      prefetch_batchs=tf.data.experimental.AUTOTUNE,
  ):
      """Input function which provides batches for train or eval.
  
      Args:
        is_training: A boolean denoting whether the input is for training.
        data_dir: The directory containing the input data.
        batch_size: The number of samples per batch.
        dtype: Data type to use for images/features
        datasets_num_private_threads: Number of private threads for tf.data.
        input_context: A `tf.distribute.InputContext` object passed in by
          `tf.distribute.Strategy`.
        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.
        dataset_cache: Whether to cache the dataset on workers.
           Typically used to improve training performance when training data is in
           remote storage and can fit into worker memory.
        filenames: Optional field for providing the file names of the TFRecords.
        prefetch_batchs: The number of batchs to prefetch.
  
      Returns:
        A dataset that can be used for iteration.
      """
      dataset = tf.data.Dataset.from_tensor_slices(filenames)
  
      if input_context:
          # TODO(swang): Shard and set shard index based on TF session.
          logging.info(
              "Sharding the dataset: input_pipeline_id=%d num_input_pipelines=%d",
              input_context.input_pipeline_id,
              input_context.num_input_pipelines,
          )
          dataset = dataset.shard(
              input_context.num_input_pipelines,
              input_context.input_pipeline_id,
          )
  
      if is_training:
          # Shuffle the input files
          dataset = dataset.shuffle(buffer_size=len(filenames))
  
      # Convert to individual records.
      # cycle_length = 10 means that up to 10 files will be read and deserialized in
      # parallel. You may want to increase this number if you have a large number of
      # CPU cores.
      dataset = dataset.interleave(
          tf.data.TFRecordDataset,
          cycle_length=10,
          num_parallel_calls=tf.data.experimental.AUTOTUNE,
      )
  
      if is_training:
          dataset = dataset.map(
              parse_example_proto_and_decode,
              num_parallel_calls=tf.data.experimental.AUTOTUNE,
          )
      dataset = dataset.take(num_images_per_epoch)
      if dataset_cache:
          # Improve training / eval performance when data is in remote storage and
          # can fit into worker memory.
          dataset = dataset.materialize()
  
      return process_record_dataset(
          dataset=dataset,
          is_training=is_training,
          batch_size=batch_size,
          num_epochs=num_epochs,
          shuffle_buffer=shuffle_buffer,
          dtype=dtype,
          datasets_num_private_threads=datasets_num_private_threads,
          drop_remainder=drop_remainder,
          tf_data_experimental_slack=tf_data_experimental_slack,
          prefetch_batchs=prefetch_batchs,
      )
  
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 5 for calling a global function with example usage: build_tf_dataset(filenames, batch_size, num_images_per_epoch, num_epochs, shuffle_buffer, is_training, dtype, datasets_num_private_threads, input_context, drop_remainder, tf_data_experimental_slack, dataset_cache, prefetch_batchs) and returns: process_record_dataset