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.
  
      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.
      """
      # Assert last dimension is same as self._embedding_dim
      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, self._w)
                   + tf.reduce_sum(self._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)
      q_latent_loss = tf.reduce_mean((quantized - tf.stop_gradient(inputs)) ** 2)
      loss = q_latent_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,}