docs: def train(logdir):
    """Run a network on the PTB training set, checkpointing the weights."""
  
    ptb_train = PTB(
        name="ptb_train",
        subset="train",
        seq_len=FLAGS.unroll_steps,
        batch_size=FLAGS.batch_size)
  
    # Connect to training set.
    data_ops = ptb_train()
    embed_layer, rnn_core, output_linear = build_modules(
        is_training=True, vocab_size=ptb_train.vocab_size)
    prediction_logits, zero_state_op = build_logits(
        data_ops, embed_layer, rnn_core, output_linear, name_prefix="train")
    data_loss = build_loss(prediction_logits, data_ops.sparse_target)
  
    # Add the KL cost.
    total_kl_cost = bbb.get_total_kl_cost()
    num_dataset_elements = FLAGS.batch_size * ptb_train.num_batches
    scaled_kl_cost = total_kl_cost / num_dataset_elements
    total_loss = tf.add(scaled_kl_cost, data_loss)
  
    # Optimize as usual.
    global_step = tf.get_variable(
        "num_weight_updates",
        initializer=tf.constant(0, dtype=tf.int32, shape=()),
        collections=[tf.GraphKeys.GLOBAL_VARIABLES, tf.GraphKeys.GLOBAL_STEP])
  
    learning_rate = tf.get_variable(
        "lr", initializer=tf.constant(FLAGS.lr_start, shape=(), dtype=tf.float32))
    learning_rate_update = learning_rate.assign(learning_rate * FLAGS.lr_decay)
    optimizer = tf.train.GradientDescentOptimizer(
        learning_rate=learning_rate)
    optimizer = GlobalNormClippingOptimizer(optimizer, clip_norm=5.0)
  
    with tf.control_dependencies([optimizer.minimize(total_loss)]):
      global_step_and_train = global_step.assign_add(1)
  
    # Connect to valid set.
    ptb_valid = PTB(
        name="ptb_valid",
        subset="valid",
        seq_len=FLAGS.unroll_steps,
        batch_size=FLAGS.batch_size)
    valid_data_ops = ptb_valid()
    valid_logits, zero_valid_state = build_logits(
        valid_data_ops, embed_layer, rnn_core, output_linear, name_prefix="valid")
    valid_loss = build_loss(valid_logits, valid_data_ops.sparse_target)
  
    # Compute metrics for the sake of monitoring training.
    predictions = tf.cast(
        tf.argmax(prediction_logits, axis=-1), tf.int32, name="pred")
    correct_prediction_mask = tf.cast(
        tf.equal(predictions, data_ops.sparse_target), tf.int32)
    accuracy = tf.reduce_mean(
        tf.cast(correct_prediction_mask, tf.float32), name="acc")
    error_rate = tf.subtract(1.0, accuracy, name="err")
    label_probs = tf.nn.softmax(prediction_logits, dim=-1)
    predictive_entropy = tf.reduce_mean(
        label_probs * tf.log(label_probs + 1e-12) * -1.0)
  
    # Create tf.summary ops.
    log_ops_to_run = {
        "scalar": collections.OrderedDict([
            ("task_loss", data_loss),
            ("train_err_rate", error_rate),
            ("pred_entropy", predictive_entropy),
            ("learning_rate", learning_rate),
            ("elbo_loss", total_loss),
            ("kl_cost", total_kl_cost),
            ("scaled_kl_cost", scaled_kl_cost),
        ]),
        "text": collections.OrderedDict([
            ("labels", ptb_train.to_string_tensor(data_ops.sparse_target)),
            ("predictions", ptb_train.to_string_tensor(predictions))
        ])
    }
  
    for name, tensor in log_ops_to_run["scalar"].items():
      tf.summary.scalar(os.path.join("train", name), tensor)
  
    # The remaining logic runs the training loop and logging.
    summary_writer = tf.summary.FileWriterCache.get(logdir=logdir)
    tf.logging.info(
        "Beginning training for {} epochs, each with {} batches.".format(
            FLAGS.num_training_epochs, ptb_train.num_batches))
    with tf.train.MonitoredTrainingSession(
        is_chief=True, checkpoint_dir=logdir, save_summaries_secs=10) as sess:
      num_updates_v = _run_session_with_no_hooks(sess, global_step)
      epoch_idx_start, step_idx_start = divmod(
          num_updates_v, ptb_train.num_batches)
      tf.logging.info("On start, epoch: {}\t step: {}".format(
          epoch_idx_start, step_idx_start))
      for epoch_idx in range(epoch_idx_start, FLAGS.num_training_epochs):
        tf.logging.info("Beginning Epoch {}/{}".format(
            epoch_idx, FLAGS.num_training_epochs))
        tf.logging.info(
            ("Beginning by evaluating on the validation set, which has "
             "{} batches.".format(ptb_valid.num_batches)))
        valid_cost = 0
        valid_steps = 0
        _run_session_with_no_hooks(sess, zero_valid_state)
        for _ in range(ptb_valid.num_batches):
          valid_cost_v, num_updates_v = _run_session_with_no_hooks(
              sess, [valid_loss, global_step])
          valid_cost += valid_cost_v
          valid_steps += FLAGS.unroll_steps
        tf.logging.info("Validation set perplexity: {}".format(
            np.exp(valid_cost / valid_steps)))
        summary = tf.summary.Summary()
        summary.value.add(
            tag="valid/word_level_perplexity",
            simple_value=np.exp(valid_cost / valid_steps))
        summary_writer.add_summary(summary, num_updates_v)
  
        # Run a training epoch.
        epoch_cost = 0
        epoch_steps = 0
        for batch_idx in range(step_idx_start, ptb_train.num_batches):
          scalars_res, num_updates_v = sess.run(
              [log_ops_to_run["scalar"], global_step_and_train])
          epoch_cost += scalars_res["task_loss"]
          epoch_steps += FLAGS.unroll_steps
          if (batch_idx - 1) % FLAGS.print_every_batches == 0:
            summary = tf.summary.Summary()
            summary.value.add(
                tag="train/word_level_perplexity",
                simple_value=np.exp(epoch_cost / epoch_steps))
            summary_writer.add_summary(summary, num_updates_v)
            scalars_res, strings_res = _run_session_with_no_hooks(
                sess, [log_ops_to_run["scalar"], log_ops_to_run["text"]])
            tf.logging.info("Num weight updates: {}".format(num_updates_v))
            for name, result in scalars_res.items():
              tf.logging.info("{}: {}".format(name, result))
            for name, result in strings_res.items():
              tf.logging.info("{}: {}".format(name, result))
  
        word_level_perplexity = np.exp(epoch_cost / epoch_steps)
        tf.logging.info(
            "Train Perplexity after Epoch {}: {}".format(
                epoch_idx, word_level_perplexity))
  
        end_of_epoch_fetches = [zero_state_op]
        if epoch_idx >= FLAGS.high_lr_epochs:
          end_of_epoch_fetches.append(learning_rate_update)
        _run_session_with_no_hooks(sess, end_of_epoch_fetches)
  
    tf.logging.info("Done training. Thanks for your time.")