code: def embedding_lookup(self, x, means):
      
      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
docs: 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.