code: def vq_nearest_neighbor(x, means,
                          soft_em=False, num_samples=10, temperature=None):
    
    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
docs: Find the nearest element in means to elements in x.