class_docstr: Discrete bottleneck class.
class_docstr_tok: ["Discrete","bottleneck","class","."]
docstr: 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.
docstr_tok: ["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","."]
code: 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,
      }
  
file_path: tensor2tensor/tensor2tensor/layers/vq_discrete.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from functools import partial
  from tensor2tensor.layers import common_layers
  import tensorflow.compat.v1
  from tensorflow.python.training import moving_averages
label_desc: Use image node_id 8 for calling the DiscreteBottleneck obj's underlying member method code with example usage: obj.discrete_bottleneck(x) and returns: dict