code: def vq_nearest_neighbor(
      x, means, soft_em=False, num_samples=10, temperature=None
  ):
      """Find the nearest element in means to elements in x."""
      bottleneck_size = common_layers.shape_list(means)[0]
      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 soft_em:
          x_means_idx = tf.multinomial(-dist, num_samples=num_samples)
          x_means_hot = tf.one_hot(
              x_means_idx, depth=common_layers.shape_list(means)[0]
          )
          x_means_hot = tf.reduce_mean(x_means_hot, axis=1)
      else:
          if temperature is None:
              x_means_idx = tf.argmax(-dist, axis=-1)
          else:
              x_means_idx = tf.multinomial(-dist / temperature, 1)
              x_means_idx = tf.squeeze(x_means_idx, axis=-1)
          if (
              common_layers.should_generate_summaries()
              and not common_layers.is_xla_compiled()
          ):
              tf.summary.histogram(
                  "means_idx", tf.reshape(x_means_idx, [-1])
              )
          x_means_hot = tf.one_hot(x_means_idx, bottleneck_size)
      x_means_hot_flat = tf.reshape(x_means_hot, [-1, bottleneck_size])
      x_means = tf.matmul(x_means_hot_flat, means)
      e_loss = tf.reduce_mean(
          tf.squared_difference(x, tf.stop_gradient(x_means))
      )
      return x_means_hot, e_loss, dist
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from functools import partial
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.layers import common_image_attention
  from tensor2tensor.layers import common_layers
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tensorflow_probability
  from tensorflow.python.training import moving_averages
label_desc: Use image node_id 16 for calling a global function with example usage: vq_nearest_neighbor(x, means, soft_em, num_samples, temperature) and returns: x_means_hot,  e_loss,  dist