docs: 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,}