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).

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