import tensorflow as tf

  perturbs = [tf.stop_gradient(d) for d in perturbs]

  perturbs = [_scale_l2(_mask_by_length(d, length), perturb_norm_length) for d in perturbs]
  vadv_logits = logits_from_embedding_fn([emb + d for (emb, d) in zip(embedded, perturbs)])
  return _kl_divergence_with_logits(logits, vadv_logits, weights, num_classes)


def _mask_by_length(t, length):
  maxlen = t.get_shape().as_list()[1]
  mask = tf.sequence_mask(length, maxlen=maxlen)
  mask = tf.expand_dims(tf.cast(mask, tf.float32), -1)
  return t * mask


def _scale_l2(x, norm_length):
  alpha = tf.reduce_max(tf.abs(x), (1, 2), keep_dims=True) + 1e-12
  l2_norm = alpha * tf.sqrt(tf.reduce_sum(tf.pow(x / alpha, 2), (1, 2), keep_dims=True) + 1e-6)
  x_unit = x / l2_norm
  return norm_length * x_unit


def _end_of_seq_mask(tokens, vocab_size):
  """Generate a mask for the EOS token (1.0 on EOS, 0.0 otherwise).

  Args:
      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:
