Compute the loss of large vocab tensors using a VQAE codebook.

  Args:
    x: Tensor of inputs to be quantized to nearest code
    targets: Tensor of target indices to target codes
    codebook_size: Size of quantization codebook
    beta: scalar float for moving averages
    decay: scalar float for moving averages
    epsilon: scalar float for moving averages
    soft_em: boolean, whether to apply a soft sampling procedure
    num_samples: if soft_em, number of samples to take
    temperature: temperature if we want to sample nearest neighbors or None
    do_update: whether to update the means; True by default, can be a Tensor

  Returns:
    discrete_x: one-hot Tensor indicating which codebook element is closest to x
    x_means: Tensor, on the forward pass: closest codebook element to x, on the
      backwards pass: soft convex-combination of codebook elements by proximity
      to x
    target_means: the codebook elements corresponding to the targets
    code_loss: loss driving x closer to its nearest codebook element
    targets_loss: cross-entropy loss driving x closer to code corresponding to
      target

def vq_loss(x,
            targets,
            codebook_size,
            beta=0.25,
            decay=0.999,
            epsilon=1e-5,
            soft_em=False,
            num_samples=10,
            temperature=None,
            do_update=True):
  """Compute the loss of large vocab tensors using a VQAE codebook.

  Args:
    x: Tensor of inputs to be quantized to nearest code
    targets: Tensor of target indices to target codes
    codebook_size: Size of quantization codebook
    beta: scalar float for moving averages
    decay: scalar float for moving averages
    epsilon: scalar float for moving averages
    soft_em: boolean, whether to apply a soft sampling procedure
    num_samples: if soft_em, number of samples to take
    temperature: temperature if we want to sample nearest neighbors or None
    do_update: whether to update the means; True by default, can be a Tensor

  Returns:
    discrete_x: one-hot Tensor indicating which codebook element is closest to x
    x_means: Tensor, on the forward pass: closest codebook element to x, on the
      backwards pass: soft convex-combination of codebook elements by proximity
      to x
    target_means: the codebook elements corresponding to the targets
    code_loss: loss driving x closer to its nearest codebook element
    targets_loss: cross-entropy loss driving x closer to code corresponding to
      target
  """
  x_shape = common_layers.shape_list(x)
  target_shape = common_layers.shape_list(targets)
  hidden_size = x_shape[-1]
  means, _, _ = get_vq_codebook(codebook_size, hidden_size)
  x = tf.reshape(x, [-1, hidden_size])
  targets = tf.reshape(targets, [-1])
  one_hot_targets = tf.one_hot(targets, codebook_size)
  target_means = tf.matmul(one_hot_targets, means)

  discrete_x, code_loss, distances = vq_body(
      x,
      codebook_size,
      beta=beta,
      decay=decay,
      epsilon=epsilon,
      soft_em=soft_em,
      num_samples=num_samples,
      temperature=temperature,
      do_update=do_update)

  logits = -distances
  targets_loss = tf.losses.sparse_softmax_cross_entropy(
      logits=logits, labels=targets)
  targets_loss = tf.reduce_mean(targets_loss)

  x_means = tf.matmul(discrete_x, means)
  x_means = x + tf.stop_gradient(x_means - x)

  discrete_x = tf.reshape(discrete_x, x_shape[:-1] + [codebook_size])
  target_means = tf.reshape(target_means, target_shape + [hidden_size])
  return discrete_x, x_means, target_means, code_loss, targets_loss