code: def vq_discrete_bottleneck(x, hparams):
    
    tf.logging.info("Using EMA with beta = {}".format(hparams.beta))
    bottleneck_size = 2**hparams.bottleneck_bits
    x_shape = common_layers.shape_list(x)
    x = tf.reshape(x, [-1, hparams.hidden_size])
    x_means_hot, e_loss = vq_nearest_neighbor(
        x, hparams)
    means, ema_means, ema_count = (hparams.means, hparams.ema_means,
                                   hparams.ema_count)
  
    # Update the ema variables
    updated_ema_count = moving_averages.assign_moving_average(
        ema_count,
        tf.reduce_sum(x_means_hot, axis=0),
        hparams.decay,
        zero_debias=False)
  
    dw = tf.matmul(x_means_hot, x, transpose_a=True)
    updated_ema_means = moving_averages.assign_moving_average(
        ema_means, dw, hparams.decay, zero_debias=False)
    n = tf.reduce_sum(updated_ema_count, axis=-1, keepdims=True)
    updated_ema_count = (
        (updated_ema_count + hparams.epsilon) /
        (n + bottleneck_size * hparams.epsilon) * n)
    # pylint: disable=g-no-augmented-assignment
    updated_ema_means = updated_ema_means / tf.expand_dims(
        updated_ema_count, axis=-1)
    # pylint: enable=g-no-augmented-assignment
    with tf.control_dependencies([e_loss]):
      update_means = tf.assign(means, updated_ema_means)
      with tf.control_dependencies([update_means]):
        loss = hparams.beta * e_loss
  
    discrete = tf.reshape(x_means_hot, x_shape[:-1] + [bottleneck_size])
    return discrete, loss
docs: Simple vector quantized discrete bottleneck.