code: def nearest_neighbor(
      x,
      means,
      block_v_size,
      random_top_k=1,
      soft_em=False,
      num_samples=1,
      sum_over_latents=False,
      summary=True,
  ):
      """Find the nearest element in means to elements in x.
  
      Args:
        x: Continuous encodings of shape [batch_size, latent_dim, num_blocks,
          block_dim].
        means: Embedding table of shape [num_blocks, block_v_size, block_dim].
        block_v_size: Number of table entries per block.
        random_top_k: Noisy top-k if this is bigger than 1.
        soft_em: If True then use soft EM rather than hard EM.
        num_samples: Number of samples to take in soft EM.
        sum_over_latents: Whether to sum over non-batch dimensions when calculating
          negative entropy loss. Used only when doing soft EM.
        summary: If True then record summary histogram of entropies.
  
      Returns:
        Tensor with nearest element in mean encoded in one-hot notation
        and distances.
      """
      (
          batch_size,
          latent_dim,
          num_blocks,
          block_dim,
      ) = common_layers.shape_list(x)
      x = tf.reshape(
          x, [batch_size * latent_dim, num_blocks, block_dim]
      )
      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
      )
  
      # computing cluster probabilities
      if soft_em:
          num_blocks = common_layers.shape_list(dist)[1]
          nearest_idx = tf.stack(
              [
                  tf.multinomial(
                      -dist[:, i, :], num_samples=num_samples
                  )
                  for i in range(num_blocks)
              ],
              axis=1,
          )
          nearest_hot = tf.one_hot(nearest_idx, depth=block_v_size)
          neg_q_entropy = tf.reduce_sum(
              nearest_hot * tf.expand_dims(tf.nn.log_softmax(-dist), 2),
              axis=2,
          )
          if sum_over_latents:
              neg_q_entropy = tf.reduce_sum(neg_q_entropy, [1, 2])
          neg_q_entropy = tf.reduce_mean(neg_q_entropy, axis=0)
          nearest_hot = tf.reduce_mean(nearest_hot, axis=-2)
          if summary:
              tf.summary.histogram(
                  "neg_q_entropy", tf.reshape(neg_q_entropy, [-1])
              )
      else:
          neg_q_entropy = 0.0
          if random_top_k > 1:
              _, top_k_idx = tf.nn.top_k(-dist, k=random_top_k)
              nearest_idx = tf.gather(
                  top_k_idx,
                  tf.random_uniform(
                      [1],
                      minval=0,
                      maxval=random_top_k - 1,
                      dtype=tf.int32,
                  ),
                  axis=-1,
              )
          else:
              nearest_idx = tf.argmax(-dist, axis=-1)
          nearest_hot = tf.one_hot(nearest_idx, block_v_size)
      return nearest_hot, neg_q_entropy
  
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 3 for calling a global function with example usage: nearest_neighbor(x, means, block_v_size, random_top_k, soft_em, num_samples, sum_over_latents, summary) and returns: nearest_hot,  neg_q_entropy