def loss_with_update():
    """Update the ema variables and return loss triggering the update."""
    updated_ema_count = moving_averages.assign_moving_average(
        ema_count,
        tf.reduce_sum(tf.reshape(x_means_hot, shape=[-1, codebook_size]),
                      axis=0),
        decay,
        zero_debias=False)

    dw = tf.matmul(x_means_hot, x, transpose_a=True)
    updated_ema_means = tf.identity(
        moving_averages.assign_moving_average(
            ema_means, dw, decay, zero_debias=False))
    n = tf.reduce_sum(updated_ema_count, axis=-1, keepdims=True)
    updated_ema_count = (
        (updated_ema_count + epsilon) / (n + codebook_size * epsilon) * n)
    updated_ema_means /= tf.expand_dims(updated_ema_count, axis=-1)
    with tf.control_dependencies([e_loss]):
      update_means = means.assign(updated_ema_means)
      with tf.control_dependencies([update_means]):
        return beta * e_loss
Update the ema variables and return loss triggering the update.
tensor2tensor/layers/discretization.py