def model_function(features, labels, mode):
  """A generic model, satisfying the tf.Estimator API.

  Args:
    features: Dictionary of input tensors.
      Key "seqs":
        For the pattern task: <tf.float32>[batch_size, seq_len, pattern_size]
        Sequences of patterns, where `seq_len` is the length of a sequence
        (including the query pattern) and `pattern_size` is the dimensionality
        of each pattern.
        For the symbolic task: <tf.int32>[batch_size, seq_len]
        Sequences of encoded kv pairs.
      Key "targets": same as `labels` arg.
    labels:
        For the pattern task: <tf.int32>[batch_size, pattern_size]
        The correct pattern to retrieve for the degraded query.
        For the symbolic task: <tf.int32>[batch_size]
        The correct value to retrieve for the given key.
    mode: One of the keys from tf.estimator.ModeKeys

  Returns:
    final_out: retrieved pattern for the degraded query.

  Raises:
    NameError: If FLAGS.model_name is not one of "vanilla_rnn", "vanilla_lstm",
    "rnn_attention", "rnn_fast_weights", "rnn_differentiable_plasticity," or if
    FLAGS.task_name is not "pattern" or "symbolic".
  """
  seqs = features["seqs"]

  if FLAGS.task_name == "symbolic":
    # seqs: <tf.int32>[FLAGS.batch_size, len(seqs), SYMBOLIC_VOCAB_SIZE]
    #   where len(seqs) == 2 * FLAGS.num_sets_per_sequence as each set is a
    #   key-value pair to put into memory.
    # labels: <tf.int32>[FLAGS.batch_size, SYMBOLIC_VOCAB_SIZE]
    seqs = tf.one_hot(seqs, SYMBOLIC_VOCAB_SIZE)
    labels = tf.one_hot(labels, SYMBOLIC_VOCAB_SIZE)

  # Forward pass through model
  # preds: <float32>[FLAGS.batch_size, FLAGS.pattern_size]
  if FLAGS.model_name == "vanilla_rnn":
    preds = vanilla_rnn(seqs, hidden_size=FLAGS.hidden_size)
  elif FLAGS.model_name == "vanilla_lstm":
    preds = vanilla_lstm(seqs, hidden_size=FLAGS.hidden_size)
  elif FLAGS.model_name == "rnn_attention":
    preds = rnn_attention(
        seqs, batch_size=FLAGS.batch_size, hidden_size=FLAGS.hidden_size)
  elif FLAGS.model_name == "rnn_fast_weights":
    preds = rnn_fast_weights(
        seqs,
        batch_size=FLAGS.batch_size,
        hidden_size=FLAGS.hidden_size,
        fast_steps=FLAGS.fast_steps,
        fast_decay_rate=FLAGS.fast_decay_rate,
        fast_lr=FLAGS.fast_lr,
    )
  elif FLAGS.model_name == "rnn_differentiable_plasticity":
    preds = rnn_differentiable_plasticity(
        seqs,
        batch_size=FLAGS.batch_size,
        hidden_size=FLAGS.hidden_size,
        fast_steps=FLAGS.fast_steps,
        fast_lr_fixed=FLAGS.fast_lr,
        use_oja_rule=FLAGS.use_oja,
        update_mem_with_prev_timestep=FLAGS.update_mem_with_prev_timestep,
        learn_fast_lr=FLAGS.learn_fast_lr,
        learn_plasticity_coeffs=FLAGS.learn_plasticity_coeffs,
    )
  elif FLAGS.model_name == "rnn_explicit_mem":
    preds = rnn_explicit_mem(
        seqs, batch_size=FLAGS.batch_size, hidden_size=FLAGS.hidden_size)
  else:
    raise NameError("Model %s not found" % FLAGS.model_name)

  # Calculate loss based on the last output (retrieved for the query)
  if FLAGS.task_name == "pattern":
    loss = tf.losses.mean_squared_error(preds, labels)
  elif FLAGS.task_name == "symbolic":
    loss = tf.reduce_mean(
        tf.nn.softmax_cross_entropy_with_logits(logits=preds, labels=labels))
  else:
    raise NameError("Task name %s not found" % FLAGS.task_name)

  lr = tf.train.exponential_decay(
      FLAGS.learning_rate,
      tf.train.get_global_step(),
      FLAGS.lr_decay_step,
      FLAGS.lr_decay_rate,
      staircase=True)
  tf.summary.scalar("lr", lr)

  optimizer = tf.train.AdamOptimizer(learning_rate=lr)
  grads, variables = zip(
      *optimizer.compute_gradients(loss, tf.trainable_variables()))
  grads, _ = tf.clip_by_global_norm(grads, FLAGS.max_gradient_norm)
  train_op = optimizer.apply_gradients(
      zip(grads, variables), global_step=tf.train.get_global_step())

  logging_hook = tf.train.LoggingTensorHook({"loss": loss}, every_n_iter=1000)

  if FLAGS.task_name == "pattern":
    eval_metric_ops = {
        "hamming": hamming_loss(preds, labels, sign=True),
        # Naive baseline where we just guess the query
        "naive_guess_query": tf.metrics.mean_squared_error(
            seqs[:, -1, :], tf.cast(labels, seqs.dtype))
    }
  elif FLAGS.task_name == "symbolic":
    eval_metric_ops = {
        "accuracy": tf.metrics.accuracy(predictions=tf.argmax(preds, 1),
                                        labels=tf.argmax(labels, 1))
    }

  return tf_estimator.EstimatorSpec(
      mode=mode,
      loss=loss,
      predictions=preds,
      train_op=train_op,
      eval_metric_ops=eval_metric_ops,
      training_hooks=[logging_hook],
  )
A generic model, satisfying the tf.Estimator API.

  Args:
    features: Dictionary of input tensors.
      Key "seqs":
        For the pattern task: <tf.float32>[batch_size, seq_len, pattern_size]
        Sequences of patterns, where `seq_len` is the length of a sequence
        (including the query pattern) and `pattern_size` is the dimensionality
        of each pattern.
        For the symbolic task: <tf.int32>[batch_size, seq_len]
        Sequences of encoded kv pairs.
      Key "targets": same as `labels` arg.
    labels:
        For the pattern task: <tf.int32>[batch_size, pattern_size]
        The correct pattern to retrieve for the degraded query.
        For the symbolic task: <tf.int32>[batch_size]
        The correct value to retrieve for the given key.
    mode: One of the keys from tf.estimator.ModeKeys

  Returns:
    final_out: retrieved pattern for the degraded query.

  Raises:
    NameError: If FLAGS.model_name is not one of "vanilla_rnn", "vanilla_lstm",
    "rnn_attention", "rnn_fast_weights", "rnn_differentiable_plasticity," or if
    FLAGS.task_name is not "pattern" or "symbolic".
  
language/labs/memory/run_models.py