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.

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.
    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