import tensorflow as tf

  Args:
      q_logits: logits for 1st argument of KL divergence shape
            [num_timesteps * batch_size, num_classes] if num_classes > 2, and
            [num_timesteps * batch_size] if num_classes == 2.
      p_logits: logits for 2nd argument of KL divergence with same shape q_logits.
      weights: 1-D `float` tensor with shape [num_timesteps * batch_size].
           Elements should be 1.0 only on end of sequences
      num_classes: a `int`, number of training classes

  Returns:
      a `float` `scalar`, KL divergence.
  """
  if num_classes == 2:
    q = tf.nn.sigmoid(q_logits)
    p = tf.nn.sigmoid(p_logits)
    kl = (-tf.nn.sigmoid_cross_entropy_with_logits(logits=q_logits, labels=q) +
          f.nn.sigmoid_cross_entropy_with_logits(logits=p_logits, labels=q))

  else:
    q = tf.nn.softmax(q_logits)
    p = tf.nn.softmax(p_logits)
    kl = tf.reduce_sum(q * (tf.log(q) - tf.log(p)), 1)

  num_labels = tf.reduce_sum(weights)
  num_labels = tf.where(tf.equal(num_labels, 0.), 1., num_labels)

  kl.get_shape().assert_has_rank(2)
  weights.get_shape().assert_has_rank(1)
