def input_function(is_train):
  """Pulls in pattern dataset.

  Args:
    is_train: A boolean indicating whether we are constructing the dataset for
    the train set or the eval set.

  Returns:
    dataset: A tf.data.Dataset object containing the features and labels.

  Raises:
    NameError: If FLAGS.task_name is not "pattern" or "symbolic".
  """
  if FLAGS.task_name == "pattern":
    dataset = synthetic_dataset.get_pattern_dataset(
        FLAGS.num_examples,
        FLAGS.num_sets_per_sequence,
        FLAGS.pattern_size,
        selective=FLAGS.selective_task,
        num_patterns_store=FLAGS.num_patterns_store)
  elif FLAGS.task_name == "symbolic":
    dataset = synthetic_dataset.get_symbolic_dataset(
        is_train,
        FLAGS.num_examples,
        FLAGS.num_sets_per_sequence)
  else:
    raise NameError("Task %s not found" % FLAGS.task_name)

  dataset = (
      dataset.repeat(FLAGS.num_epochs).shuffle(buffer_size=1000)
      .batch(FLAGS.batch_size, drop_remainder=True))

  # Estimator expects a tuple.
  dataset = dataset.map(lambda d: (d, d["targets"]))

  return dataset
Pulls in pattern dataset.

  Args:
    is_train: A boolean indicating whether we are constructing the dataset for
    the train set or the eval set.

  Returns:
    dataset: A tf.data.Dataset object containing the features and labels.

  Raises:
    NameError: If FLAGS.task_name is not "pattern" or "symbolic".
  
language/labs/memory/run_models.py