def METHOD_NAME(input_dir, mode, batch_size):
  """Reads data to be inputted in TF experiment from DataFlow output.
  Processed data, stored as TFRecord.
  Args:
    input_dir: Directory to read data from (output of DataFlow job).
    mode: Specifies the type of dataset (train, validation, test).
    batch_size: Batch size used to read data.
  Raises:
    ValueError: Expected one of the following: `tf.estimator.ModeKeys.TRAIN`,
    `tf.estimator.ModeKeys.EVAL`, `tf.estimator.ModeKeys.INFER`.
  Returns:
    Input function.
  """
  def _parser(data):
    features = {c: data[c] for c in constants.FEATURE_COLUMNS}
    key = constants.KEY_COLUMN
    features.update({key: data[key]})
    labels = data[constants.LABEL_COLUMN]
    return features, labels
  if mode == tf.estimator.ModeKeys.TRAIN:
    d = DatasetType.TRAIN
  elif mode == tf.estimator.ModeKeys.EVAL:
    d = DatasetType.VAL
  elif mode == tf.estimator.ModeKeys.INFER:
    d = DatasetType.TEST
  else:
    raise ValueError(
        'Expected one of the following: tf.estimator.ModeKeys.TRAIN, EVAL, '
        'INFER ; got {} instead.'.format(mode))
  assert d
  prefix_input = [
      posixpath.join(input_dir,
                     '{}*'.format(constants.PATH_TRANSFORMED_DATA_SPLIT[d]))
  ]
  def _input_fn():
    """Input function that serves as input to the Experiment class.
    Returns:
      `tf.Data.dataset` object containing features and labels.
    """
    filenames = tf.train.match_filenames_once(prefix_input)
    dataset = tf.data.TFRecordDataset(filenames)
    schema = make_input_schema()
    dataset = dataset.map(
        lambda x: tf.parse_single_example(serialized=x, features=schema))
    dataset = dataset.map(_parser)
    dataset = dataset.repeat(None)
    dataset = dataset.batch(batch_size)
    return dataset
  return _input_fn