import tensorflow as tf

      Tensor with nearest element in mean encoded in one-hot notation.
    """
    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

    if self.hparams.soft_em:
      nearest_idx = tf.stack(
