code: def vq_nearest_neighbor(x, hparams):
      """Find the nearest element in means to elements in x."""
      bottleneck_size = 2**hparams.bottleneck_bits
      means = hparams.means
      x_norm_sq = tf.reduce_sum(tf.square(x), axis=-1, keepdims=True)
      means_norm_sq = tf.reduce_sum(
          tf.square(means), axis=-1, keepdims=True
      )
      scalar_prod = tf.matmul(x, means, transpose_b=True)
      dist = x_norm_sq + tf.transpose(means_norm_sq) - 2 * scalar_prod
      if hparams.bottleneck_kind == "em":
          x_means_idx = tf.multinomial(
              -dist, num_samples=hparams.num_samples
          )
          x_means_hot = tf.one_hot(x_means_idx, depth=bottleneck_size)
          x_means_hot = tf.reduce_mean(x_means_hot, axis=1)
      else:
          x_means_idx = tf.argmax(-dist, axis=-1)
          x_means_hot = tf.one_hot(x_means_idx, depth=bottleneck_size)
      x_means = tf.matmul(x_means_hot, means)
      e_loss = tf.reduce_mean(
          tf.squared_difference(x, tf.stop_gradient(x_means))
      )
      return x_means_hot, e_loss
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  from six.moves import range
  from tensor2tensor.layers import common_layers
  from tensor2tensor.models import transformer
  from tensor2tensor.utils import beam_search
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import t2t_model
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  from tensorflow.python.training import moving_averages
label_desc: Use image node_id 2 for calling a global function with example usage: vq_nearest_neighbor(x, hparams) and returns: x_means_hot,  e_loss