bc08dbfbe77554bf3458529285003f0bf92eecd3,sonnet/python/modules/nets/vqvae.py,VectorQuantizerEMA,_build,#VectorQuantizerEMA#Any#Any#,173

Before Change

    encoding_indices = tf.argmax(- distances, 1)
    encodings = tf.one_hot(encoding_indices, self._num_embeddings)
    quantized = tf.reshape(
        tf.matmul(encodings, w, transpose_b=True), tf.shape(inputs))
    e_latent_loss = tf.reduce_mean((tf.stop_gradient(quantized) - inputs) ** 2)

    if is_training:
      updated_ema_cluster_size = moving_averages.assign_moving_average(

After Change

    encoding_indices = tf.argmax(- distances, 1)
    encodings = tf.one_hot(encoding_indices, self._num_embeddings)
    encoding_indices = tf.reshape(encoding_indices, inputs.shape.as_list()[:-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(

Instances