code: 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
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from functools import partial
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.layers import common_image_attention
  from tensor2tensor.layers import common_layers
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tensorflow_probability
  from tensorflow.python.training import moving_averages
label_desc: Use image node_id 19 for calling a global function with example usage: vq_loss(x, targets, codebook_size, beta, decay, epsilon, soft_em, num_samples, temperature, do_update) and returns: discrete_x,  x_means,  target_means,  code_loss,  targets_loss