class_docstr: Discrete bottleneck class.
class_docstr_tok: ["Discrete","bottleneck","class","."]
docstr: Compute nearest neighbors and loss for training the embeddings.
  
  Args:
      x: Batch of encoder continuous latent states sliced/projected into
      shape
      [-1, num_blocks, block_dim].
      means: Embedding means.
  
  Returns:
      The nearest neighbor in one hot form, the nearest neighbor
      itself, the
      commitment loss, embedding training loss.
docstr_tok: ["Compute","nearest","neighbors","and","loss","for","training","the","embeddings",".","Args",":","x",":","Batch","of","encoder","continuous","latent","states","sliced\/projected","into","shape","[","-1",",","num_blocks",",","block_dim","]",".","means",":","Embedding","means",".","Returns",":","The","nearest","neighbor","in","one","hot","form",",","the","nearest","neighbor","itself",",","the","commitment","loss",",","embedding","training","loss","."]
code: def embedding_lookup(self, x, means):
      """Compute nearest neighbors and loss for training the embeddings.
  
      Args:
          x: Batch of encoder continuous latent states sliced/projected into
          shape
          [-1, num_blocks, block_dim].
          means: Embedding means.
  
      Returns:
          The nearest neighbor in one hot form, the nearest neighbor
          itself, the
          commitment loss, embedding training loss.
      """
      x_means_hot = self.nearest_neighbor(x, means)
      x_means_hot_flat = tf.reshape(
          x_means_hot,
          [-1, self.hparams.num_blocks, self.hparams.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])
      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
  
file_path: tensor2tensor/tensor2tensor/layers/vq_discrete.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from functools import partial
  from tensor2tensor.layers import common_layers
  import tensorflow.compat.v1
  from tensorflow.python.training import moving_averages
label_desc: Use image node_id 4 for calling the DiscreteBottleneck obj's underlying member method code with example usage: obj.embedding_lookup(x, means) and returns: x_means_hot,  x_means,  q_loss,  e_loss