Discretization bottleneck for latent variables.

    Args:
        x: Input to the discretization bottleneck.

    Returns:
        Embedding to pass to the decoder, discrete latent, loss, and the
        embedding
        function.

    Raises:
        ValueError: If projection_tensors is None for reshape_method
        project, or
        ema_count or ema_means is None if we are using ema, or unknown
        args.

def discrete_bottleneck(self, x):
    """Discretization bottleneck for latent variables.

    Args:
        x: Input to the discretization bottleneck.

    Returns:
        Embedding to pass to the decoder, discrete latent, loss, and the
        embedding
        function.

    Raises:
        ValueError: If projection_tensors is None for reshape_method
        project, or
        ema_count or ema_means is None if we are using ema, or unknown
        args.
    """
    x_reshaped = self.slice_hidden(x)
    x_means_hot = []
    x_means = 0
    loss = 0
    x_means_hot, x_means, q_loss, e_loss = self.embedding_lookup(
        x_reshaped, self.means)

    if self.hparams.ema:
      tf.logging.info("Using EMA with beta = {}".format(self.hparams.beta))
      updated_ema_count = \
          moving_averages.assign_moving_average(
              self.ema_count,
              tf.reduce_sum(
                  tf.reshape(
                      x_means_hot,
                      shape=[-1, self.hparams.num_blocks,
                             self.hparams.block_v_size]),
                  axis=0),
              self.hparams.decay,
              zero_debias=False)

      dw = tf.matmul(
          tf.transpose(x_means_hot, perm=[1, 2, 0]),
          tf.transpose(x_reshaped, perm=[1, 0, 2]))

      updated_ema_means = \
          moving_averages.assign_moving_average(
              self.ema_means, dw, self.hparams.decay,
              zero_debias=False)
      n = tf.reduce_sum(updated_ema_count, axis=-1, keep_dims=True)
      updated_ema_count = ((updated_ema_count + self.hparams.epsilon) / (
          n + 2**self.hparams.z_size * self.hparams.epsilon) * n)
      updated_ema_means = updated_ema_means / tf.expand_dims(
          updated_ema_count, axis=-1)

      with tf.control_dependencies([e_loss]):
        update_means = tf.assign(self.means, updated_ema_means)
        with tf.control_dependencies([update_means]):
          loss += self.hparams.beta * e_loss
    else:
      # Use a gradient based loss for learning the cluster centers
      loss += q_loss + self.hparams.beta * e_loss

    # Get the discrete latent representation
    x_means_idx = tf.argmax(x_means_hot, axis=-1)

    # Get the binary representation
    num_bits = int(self.hparams.z_size // self.hparams.num_blocks)
    x_means_bits = self.int_to_bit(x_means_idx, num_bits=num_bits, base=2)
    x_discrete = self.bit_to_int(
        tf.to_int32(x_means_bits), num_bits=self.hparams.z_size, base=2)

    # Reshape x_discrete
    shape_x = common_layers.shape_list(x)
    shape_discrete = shape_x[:-1]
    x_discrete = tf.reshape(x_discrete, shape_discrete)
    x_means = tf.reshape(x_means, shape=shape_x)
    h1 = x + tf.stop_gradient(x_means - x)

    h2 = tf.layers.dense(tf.nn.relu(h1), self.hparams.filter_size, name="vch2")
    res = tf.layers.dense(
        tf.nn.relu(h2), self.hparams.hidden_size, name="vcfin")
    embed_fn = partial(self.embed)
    return {
        "dense": res,
        "discrete": x_discrete,
        "loss": loss,
        "embed": embed_fn
    }