docs: 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)