code: def inputs(train, batch_size, num_epochs):
      """Reads input data num_epochs times.
      Args:
        train: Selects between the training ('train'),  validation ('vald') data and 'test' data.
        batch_size: Number of examples per returned batch.
        num_epochs: Number of times to read the input data, or 0/None to
           train forever.
      Returns:
        A tuple (images, labels), where:
        * images is a float tensor with shape [batch_size, mnist.IMAGE_PIXELS]
          in the range [-0.5, 0.5].
        * labels is an int32 tensor with shape [batch_size] with the true label,
          a number in the range [0, mnist.NUM_CLASSES).
        This function creates a one_shot_iterator, meaning that it will only iterate
        over the dataset once. On the other hand there is no special initialization
        required.
      """
      if not num_epochs:
          num_epochs = None
  
      if train == "train":
          filename = TRAIN_FILE
      elif train == "vald":
          filename = VALIDATION_FILE
      elif train == "test":
          filename = TEST_FILE
      elif train == "ref":
          filename = TRAIN_FILE
      else:
          raise ValueError(
              "Invalid argument `{}` -  must be one of: train, vald, test".format(
                  train
              )
          )
  
      with tf.name_scope("input"):
          # TFRecordDataset opens a binary file and reads one record at a time.
          # `filename` could also be a list of filenames, which will be read in order.
          dataset = tf.data.TFRecordDataset(filename)
  
          # The map transformation takes a function and applies it to every element
          # of the dataset.
          dataset = dataset.map(decode)
  
          if train == "train":
              dataset = dataset.map(augment)
  
          dataset = dataset.map(normalize)
  
          # The shuffle transformation uses a finite-sized buffer to shuffle elements
          # in memory. The parameter is the number of elements in the buffer. For
          # completely uniform shuffling, set the parameter to be the same as the
          # number of elements in the dataset.
          if train == "train" or train == "ref":
              dataset = dataset.shuffle(1000 + 3 * batch_size)
  
          if num_epochs is not None:
              dataset = dataset.repeat(num_epochs)
          elif train == "ref":
              dataset = dataset.repeat()
  
          dataset = dataset.batch(
              batch_size, drop_remainder=True
          )  # Drop remainder for consistent batch size
  
          if train == "ref":
              return dataset.make_one_shot_iterator()
  
          iterator = dataset.make_initializable_iterator()
      return iterator
  
file_path: attributionpriors/attributionpriors/mnist/mnist_reader.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import argparse
  import os.path
  import sys
  import time
  import math
  import tensorflow
  from tensorflow.examples.tutorials.mnist import mnist
label_desc: Use image node_id 4 for calling a global function with example usage: inputs(train, batch_size, num_epochs) and returns: iterator, dataset