from tensorflow.contrib.eager.python import tfe

def loss(labels, predictions):
  """Computes mean squared loss."""
  return tf.reduce_mean(tf.squared_difference(predictions, labels))


def test(model, eval_data):
  """Computes the average loss on eval_data, which should be a Dataset."""
  avg_loss = tfe.metrics.Mean("loss")
  for (labels, chars, sequence_length) in tfe.Iterator(eval_data):
    predictions = model((chars, sequence_length), training=False)
    avg_loss(loss(labels, predictions))
  print("eval/loss: %.6f\n" % avg_loss.result())
  with tf.contrib.summary.always_record_summaries():
    tf.contrib.summary.scalar("loss", avg_loss.result())
