class_docstr: Discrete bottleneck class.
class_docstr_tok: ["Discrete","bottleneck","class","."]
docstr: Find the nearest element in means to elements in x.
  
  Args:
      x: Batch of encoder continuous latent states sliced/projected into
         shape [-1, num_blocks, block_dim].
      means: Embedding means of shape.
  
  Returns:
    Tensor with nearest element in mean encoded in one-hot notation.
docstr_tok: ["Find","the","nearest","element","in","means","to","elements","in","x.","Args",":","x",":","Batch","of","encoder","continuous","latent","states","sliced\/projected","into","shape","[","-1",",","num_blocks",",","block_dim","]",".","means",":","Embedding","means","of","shape",".","Returns",":","Tensor","with","nearest","element","in","mean","encoded","in","one-hot","notation","."]
code: def nearest_neighbor(self, x, means):
      """Find the nearest element in means to elements in x.
  
      Args:
          x: Batch of encoder continuous latent states sliced/projected into
             shape [-1, num_blocks, block_dim].
          means: Embedding means of shape.
  
      Returns:
        Tensor with nearest element in mean encoded in one-hot notation.
      """
      x_norm_sq = tf.reduce_sum(tf.square(x), axis=-1, keep_dims=True)
      means_norm_sq = tf.reduce_sum(
          tf.square(means), axis=-1, keep_dims=True
      )
      scalar_prod = tf.matmul(
          tf.transpose(x, perm=[1, 0, 2]),
          tf.transpose(means, perm=[0, 2, 1]),
      )
      scalar_prod = tf.transpose(scalar_prod, perm=[1, 0, 2])
      dist = (
          x_norm_sq
          + tf.transpose(means_norm_sq, perm=[2, 0, 1])
          - 2 * scalar_prod
      )
  
      if self.hparams.soft_em:
          nearest_idx = tf.stack(
              [
                  tf.multinomial(
                      -dist[:, i, :],
                      num_samples=self.hparams.num_samples,
                  )
                  for i in range(self.hparams.num_blocks)
              ],
              axis=1,
          )
          nearest_hot = tf.one_hot(
              nearest_idx, depth=self.hparams.block_v_size
          )
          nearest_hot = tf.reduce_mean(nearest_hot, axis=-2)
      else:
          if self.hparams.random_top_k > 1:
              _, top_k_idx = tf.nn.top_k(
                  -dist, k=self.hparams.random_top_k
              )
              nearest_idx = tf.gather(
                  top_k_idx,
                  tf.random_uniform(
                      [1],
                      minval=0,
                      maxval=self.hparams.random_top_k - 1,
                      dtype=tf.int32,
                  ),
                  axis=-1,
              )
          else:
              if self.hparams.use_scales:
                  dist /= tf.reshape(
                      self.hparams.scales,
                      [1, 1, self.hparams.moe_num_experts],
                  )
              nearest_idx = tf.argmax(-dist, axis=-1)
          nearest_hot = tf.one_hot(
              nearest_idx, self.hparams.block_v_size
          )
      return nearest_hot
  
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 3 for calling the DiscreteBottleneck obj's underlying member method code with example usage: obj.nearest_neighbor(x, means) and returns: nearest_hot