import tensorflow as tf

      tokens: 1-D integer `Tensor` [num_timesteps*batch_size]. Each element is an
          id from the vocab.
      vocab_size: a `int`, vocabular size of the problem

  Returns:
      Float 1-D `Tensor` same shape as tokens, whose values are 1.0 on the end of
          sequence and 0.0 on the others.
  """
  eos_id = vocab_size - 1
  return tf.cast(tf.equal(tokens, eos_id), tf.float32)


def _kl_divergence_with_logits(q_logits, p_logits, weights, num_classes):
  """Returns weighted KL divergence between distributions q and p.

  Args:
      q_logits: logits for 1st argument of KL divergence shape
            [num_timesteps * batch_size, num_classes] if num_classes > 2, and
