Connects the module to some inputs.

    Args:
      inputs: Tensor, final dimension must be equal to embedding_dim. All other
        leading dimensions will be flattened and treated as a large batch.
      is_training: boolean, whether this connection is to training data. When
        this is set to False, the internal moving average statistics will not be
        updated.

    Returns:
      dict containing the following keys and values:
        quantize: Tensor containing the quantized version of the input.
        loss: Tensor containing the loss to optimize.
        perplexity: Tensor containing the perplexity of the encodings.
        encodings: Tensor containing the discrete encodings, ie which element
          of the quantized space each input element was mapped to.
        encoding_indices: Tensor containing the discrete encoding indices, ie
          which element of the quantized space each input element was mapped to.

def _build(self, inputs, is_training):
    """Connects the module to some inputs.

    Args:
      inputs: Tensor, final dimension must be equal to embedding_dim. All other
        leading dimensions will be flattened and treated as a large batch.
      is_training: boolean, whether this connection is to training data. When
        this is set to False, the internal moving average statistics will not be
        updated.

    Returns:
      dict containing the following keys and values:
        quantize: Tensor containing the quantized version of the input.
        loss: Tensor containing the loss to optimize.
        perplexity: Tensor containing the perplexity of the encodings.
        encodings: Tensor containing the discrete encodings, ie which element
          of the quantized space each input element was mapped to.
        encoding_indices: Tensor containing the discrete encoding indices, ie
          which element of the quantized space each input element was mapped to.
    """
    # Ensure that the weights are read fresh for each timestep, which otherwise
    # would not be guaranteed in an RNN setup. Note that this relies on inputs
    # having a data dependency with the output of the previous timestep - if
    # this is not the case, there is no way to serialize the order of weight
    # updates within the module, so explicit external dependencies must be used.
    with tf.control_dependencies([inputs]):
      w = self._w.read_value()
    input_shape = tf.shape(inputs)
    with tf.control_dependencies([
        tf.Assert(tf.equal(input_shape[-1], self._embedding_dim),
                  [input_shape])]):
      flat_inputs = tf.reshape(inputs, [-1, self._embedding_dim])

    distances = (tf.reduce_sum(flat_inputs**2, 1, keepdims=True)
                 - 2 * tf.matmul(flat_inputs, w)
                 + tf.reduce_sum(w ** 2, 0, keepdims=True))

    encoding_indices = tf.argmax(- distances, 1)
    encodings = tf.one_hot(encoding_indices, self._num_embeddings)
    encoding_indices = tf.reshape(encoding_indices, tf.shape(inputs)[:-1])
    quantized = self.quantize(encoding_indices)
    e_latent_loss = tf.reduce_mean((tf.stop_gradient(quantized) - inputs) ** 2)

    if is_training:
      updated_ema_cluster_size = moving_averages.assign_moving_average(
          self._ema_cluster_size, tf.reduce_sum(encodings, 0), self._decay)
      dw = tf.matmul(flat_inputs, encodings, transpose_a=True)
      updated_ema_w = moving_averages.assign_moving_average(self._ema_w, dw,
                                                            self._decay)
      n = tf.reduce_sum(updated_ema_cluster_size)
      updated_ema_cluster_size = (
          (updated_ema_cluster_size + self._epsilon)
          / (n + self._num_embeddings * self._epsilon) * n)

      normalised_updated_ema_w = (
          updated_ema_w / tf.reshape(updated_ema_cluster_size, [1, -1]))
      with tf.control_dependencies([e_latent_loss]):
        update_w = tf.assign(self._w, normalised_updated_ema_w)
        with tf.control_dependencies([update_w]):
          loss = self._commitment_cost * e_latent_loss

    else:
      loss = self._commitment_cost * e_latent_loss
    quantized = inputs + tf.stop_gradient(quantized - inputs)
    avg_probs = tf.reduce_mean(encodings, 0)
    perplexity = tf.exp(- tf.reduce_sum(avg_probs * tf.log(avg_probs + 1e-10)))

    return {'quantize': quantized,
            'loss': loss,
            'perplexity': perplexity,
            'encodings': encodings,
            'encoding_indices': encoding_indices,}