Creates a basic MNIST model using Sonnet, then trains and evaluates it.

def train_and_eval(train_batch_size, test_batch_size, num_hidden, learning_rate,
                   num_train_steps, report_every, test_every):
  """Creates a basic MNIST model using Sonnet, then trains and evaluates it."""

  data_dict = dataset_mnist.get_data("mnist", train_batch_size, test_batch_size)
  train_data = data_dict["train_iterator"]
  test_data = data_dict["test_iterator"]

  # Sonnet separates the configuration of a model from its attachment into the
  # graph. Here we configure the shape of the model, but this call does not
  # place any ops into the graph.
  mlp = snt.nets.MLP([num_hidden, data_dict["num_classes"]])

  train_images, train_labels = train_data.get_next()
  test_images, test_labels = test_data.get_next()

  # Flatten images to pass to model.
  train_images = snt.BatchFlatten()(train_images)
  test_images = snt.BatchFlatten()(test_images)

  # Call our model, which creates it in the graph. Our build function
  # is parameterized by the source of images, and here we connect the model to
  # the training images.
  train_logits = mlp(train_images)

  # Training loss and optimizer.
  loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=train_labels, logits=train_logits)
  loss_avg = tf.reduce_mean(loss)
  optimizer = tf.train.GradientDescentOptimizer(learning_rate)
  optimizer_step = optimizer.minimize(loss_avg)

  # As before, we make a second instance of our model in the graph, which shares
  # its parameters with the first instance of the model. The Sonnet Module code
  # takes care of the variable sharing for us: because we are calling the same
  # instance of Model, we will automatically reference the same, shared
  # variables.
  test_logits = mlp(test_images)
  test_classes = tf.nn.softmax(test_logits)
  test_correct = tf.nn.in_top_k(test_classes, test_labels, k=1)

  with tf.train.SingularMonitoredSession() as sess:

    for step_idx in range(num_train_steps):
      current_loss, _ = sess.run([loss_avg, optimizer_step])
      if step_idx % report_every == 0:
        tf.logging.info("Step: %4d of %d - loss: %.02f.",
                        step_idx + 1, num_train_steps, current_loss)
      if step_idx % test_every == 0:
        sess.run(test_data.initializer)
        current_correct = sess.run(test_correct)
        correct_count = np.count_nonzero(current_correct)
        tf.logging.info("Test: %d of %d correct.",
                        correct_count, test_batch_size)