docs: def build_and_train(iterations, log_stride, test=False):
    """Construct the data, model, loss and optimizer then train."""
  
    # Test mode settings.
    batch_size = 2 if test else FLAGS.batch_size
    num_mems = 2 if test else FLAGS.num_mems
    num_heads = 1 if test else FLAGS.num_mems
    num_blocks = 1 if test else FLAGS.num_mems
    head_size = 4 if test else FLAGS.head_size
    num_objects = 2 if test else FLAGS.num_objects
    num_features = 4 if test else FLAGS.num_features
    mlp_size = (20,) if test else (256, 256, 256, 256)
  
    with tf.Graph().as_default():
      t0 = time.time()
  
      # Initialize the dataset.
      dataset = dataset_nth_farthest.NthFarthest(
          batch_size, num_objects, num_features)
  
      # Create the model.
      core = snt.RelationalMemory(
          mem_slots=num_mems,
          head_size=head_size,
          num_heads=num_heads,
          num_blocks=num_blocks,
          gate_style=FLAGS.gate_style)
  
      final_mlp = snt.nets.MLP(
          output_sizes=mlp_size,
          activate_final=True)
  
      model = SequenceModel(
          core=core,
          target_size=num_objects,
          final_mlp=final_mlp)
  
      tf.logging.info("Instantiated models ({:3f})".format(time.time() - t0))
  
      # Get train and test data.
      inputs_train, labels_train = dataset.get_batch()
      inputs_test, labels_test = dataset.get_batch()
  
      # Define target accuracy.
      def compute_accuracy(logits, targets, name="accuracy"):
        correct_pred = tf.cast(
            tf.equal(tf.cast(targets, tf.int64), tf.argmax(logits, 1)),
            tf.float32)
        return tf.reduce_mean(correct_pred, name=name)
  
      # Define the loss & accuracy.
      def loss_fn(inputs, labels):
        """Creates the loss and the exports."""
        logits = model(inputs)
        labels = tf.cast(labels, tf.int32)
        loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
            logits=logits, labels=labels))
        accuracy = compute_accuracy(logits, labels)
        return loss, accuracy
  
      # Get training step counter.
      global_step = tf.get_variable(
          name="global_step",
          shape=[],
          dtype=tf.int64,
          initializer=tf.zeros_initializer(),
          trainable=False,
          collections=[tf.GraphKeys.GLOBAL_VARIABLES, tf.GraphKeys.GLOBAL_STEP])
  
      # Create the optimizer.
      learning_rate_op = tf.reduce_max([
          tf.train.exponential_decay(
              FLAGS.learning_rate,
              global_step,
              decay_steps=FLAGS.epochs // 100,
              decay_rate=0.9,
              staircase=False),
          FLAGS.min_learning_rate
      ])
      optimizer = tf.train.AdamOptimizer(learning_rate_op)
      train_loss, _ = loss_fn(inputs_train, labels_train)
      step_op = optimizer.minimize(train_loss, global_step=global_step)
  
      # Compute test accuracy
      logits_test = model(inputs_test)
      labels_test = tf.cast(labels_test, tf.int32)
      test_acc = compute_accuracy(logits_test, labels_test)
  
      tf.logging.info("Created losses and optimizers ({:3f})".format(
          time.time() - t0))
  
      # Begin Training.
      t0 = time.time()
      train_losses = []
      steps = []
      test_accs = []
      tf.logging.info("Starting training ({:3f})".format(time.time() - t0))
      with tf.train.SingularMonitoredSession() as sess:
        for it in six.moves.range(iterations):
          sess.run([step_op, learning_rate_op])
          if it % log_stride == 0:
            loss_v, acc_v = sess.run([train_loss, test_acc])
            elapsed = time.time() - t0
            tf.logging.info(
                "iter: {:2d}, train loss {:3f}; test acc {:3f} ({:3f})".format(
                    it, loss_v, acc_v, elapsed))
            train_losses.append(loss_v)
            steps.append(it)
          test_accs.append(acc_v)
    return steps, train_losses, test_accs