code: def embedding_lookup(
      x,
      means,
      num_blocks,
      block_v_size,
      bottleneck_kind="dvq",
      random_top_k=1,
      soft_em=False,
      num_samples=1,
      do_hard_gumbel_softmax=False,
      temperature_warmup_steps=150000,
      num_flows=0,
      approximate_gs_entropy=False,
      sum_over_latents=False,
  ):
      """Compute nearest neighbors and loss for training the embeddings via DVQ.
  
      Args:
        x: Continuous encodings of shape [batch_size, latent_dim, num_blocks,
          block_dim].
        means: Embedding table of shape [num_blocks, block_v_size, block_dim].
        num_blocks: Number of blocks in DVQ.
        block_v_size: Number of table entries per block.
        bottleneck_kind: Discrete bottleneck type.
        random_top_k: Noisy top-k if this is bigger than 1.
        soft_em: If True then use soft EM rather than hard EM.
        num_samples: Number of samples to use for soft EM.
        do_hard_gumbel_softmax: Whether to use hard or soft Gumbel-Softmax samples
          for gumbel-softmax-dvq bottleneck.
        temperature_warmup_steps: Number of steps it takes to decay temperature to
          0. Used only if bottleneck_kind is gumbel-softmax-dvq.
        num_flows: Number of inverse autoregressive flows for gumbel-softmax-dvq
          bottleneck.
        approximate_gs_entropy: Whether to approximate the Gumbel-Softmax density
          as a categorical distribution when calculating the sample entropy. Used
          only if bottleneck_kind is gumbel-softmax-dvq.
        sum_over_latents: Whether to sum over non-batch dimensions when calculating
          negative entropy loss. Used only if soft EM or when bottleneck_kind is
          gumbel-softmax-dvq.
  
      Returns:
        x_means_hot: The nearest neighbor in one hot form, with shape
          [batch_size * latent_dim, num_blocks, block_v_size].
        x_means: The nearest neighbor itself, with shape [batch_size * latent_dim,
          num_blocks, block_dim].
        q_loss: Scalar Tensor representing codebook loss.
        e_loss: Scalar Tensor representing commitment loss.
        neg_q_entropy: Scalar Tensor representing negative entropy of variational
          approximation (0 if it is deterministic).
      """
      if bottleneck_kind == "gumbel-softmax-dvq":
          (
              x_means_hot,
              neg_q_entropy,
          ) = gumbel_softmax_nearest_neighbor_dvq(
              x,
              means,
              block_v_size,
              hard=do_hard_gumbel_softmax,
              num_samples=num_samples,
              temperature_warmup_steps=temperature_warmup_steps,
              num_flows=num_flows,
              approximate_gs_entropy=approximate_gs_entropy,
              sum_over_latents=sum_over_latents,
          )
      else:
          x_means_hot, neg_q_entropy = nearest_neighbor(
              x,
              means,
              block_v_size,
              random_top_k,
              soft_em=soft_em,
              num_samples=num_samples,
              sum_over_latents=sum_over_latents,
          )
      x_means_hot_flat = tf.reshape(
          x_means_hot, [-1, num_blocks, block_v_size]
      )
      x_means = tf.matmul(
          tf.transpose(x_means_hot_flat, perm=[1, 0, 2]), means
      )
      x_means = tf.transpose(x_means, [1, 0, 2])
      (
          batch_size,
          latent_dim,
          num_blocks,
          block_dim,
      ) = common_layers.shape_list(x)
      x = tf.reshape(
          x, [batch_size * latent_dim, num_blocks, block_dim]
      )
  
      # Currently, we use the mean scaling for the commitment loss, as opposed to
      # summing across all non-batch dimensions.
      q_loss = tf.reduce_mean(
          tf.squared_difference(tf.stop_gradient(x), x_means)
      )
      e_loss = tf.reduce_mean(
          tf.squared_difference(x, tf.stop_gradient(x_means))
      )
      return x_means_hot, x_means, q_loss, e_loss, neg_q_entropy
  
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 4 for calling a global function with example usage: embedding_lookup(x, means, num_blocks, block_v_size, bottleneck_kind, random_top_k, soft_em, num_samples, do_hard_gumbel_softmax, temperature_warmup_steps, num_flows, approximate_gs_entropy, sum_over_latents) and returns: x_means_hot,  x_means,  q_loss,  e_loss,  neg_q_entropy