class_docstr: Autoencoder-augmented Transformer.
class_docstr_tok: ["Autoencoder-augmented","Transformer","."]
code: def __init__(self, *args, **kwargs):
      super(TransformerAE, self).__init__(*args, **kwargs)
      self.predict_mask = 1.0
  
      # Define bottleneck function
      self._hparams.bottleneck = functools.partial(
          discretization.discrete_bottleneck,
          hidden_size=self._hparams.hidden_size,
          z_size=self._hparams.z_size,
          filter_size=self._hparams.filter_size,
          bottleneck_kind=self._hparams.bottleneck_kind,
          num_blocks=self._hparams.num_blocks,
          num_residuals=self.hparams.num_residuals,
          reshape_method=self._hparams.reshape_method,
          beta=self._hparams.beta,
          ema=self._hparams.ema,
          epsilon=self._hparams.epsilon,
          decay=self._hparams.decay,
          random_top_k=self._hparams.random_top_k,
          soft_em=self.hparams.soft_em,
          num_samples=self.hparams.num_samples,
          softmax_k=self._hparams.softmax_k,
          temperature_warmup_steps=self._hparams.temperature_warmup_steps,
          do_hard_gumbel_softmax=self._hparams.do_hard_gumbel_softmax,
          num_flows=self._hparams.num_flows,
          approximate_gs_entropy=self._hparams.approximate_gs_entropy,
          discrete_mix=self._hparams.d_mix,
          noise_dev=self._hparams.noise_dev,
          startup_steps=self.hparams.startup_steps,
          summary=_DO_SUMMARIES,
      )
      # Set the discretization bottleneck specific things here
      if self._hparams.bottleneck_kind in ["dvq", "gumbel-softmax-dvq"]:
          z_size_per_residual = (
              self._hparams.z_size / self._hparams.num_residuals
          )
          block_dim = int(
              self._hparams.hidden_size // self._hparams.num_blocks
          )
          block_v_size = 2 ** (
              z_size_per_residual / self._hparams.num_blocks
          )
          block_v_size = int(block_v_size)
  
          if self._hparams.reshape_method == "project":
              tf.logging.info("Using projections for DVQ")
              tf.logging.info(
                  "Trainable projections = {}".format(
                      self._hparams.trainable_projections
                  )
              )
  
              projection_tensors = tf.get_variable(
                  name="projection",
                  shape=[
                      self._hparams.num_residuals,
                      self._hparams.num_blocks,
                      self._hparams.hidden_size,
                      block_dim,
                  ],
                  initializer=tf.initializers.glorot_uniform(),
                  trainable=self._hparams.trainable_projections,
              )
  
              self._hparams.bottleneck = functools.partial(
                  self._hparams.bottleneck,
                  projection_tensors=projection_tensors,
              )
          elif self._hparams.reshape_method == "slice":
              tf.logging.info("Using slices for DVQ")
          else:
              raise ValueError("Unknown reshape method")
  
          means = tf.get_variable(
              name="means",
              shape=[
                  self._hparams.num_residuals,
                  self._hparams.num_blocks,
                  block_v_size,
                  block_dim,
              ],
              initializer=tf.uniform_unit_scaling_initializer(),
          )
  
          # Create the shadow variables if we are using EMA
          ema_count = None
          ema_means = None
          if self._hparams.ema:
              ema_count = []
              for i in range(self._hparams.num_residuals):
                  ema_count_i = tf.get_variable(
                      "ema_count_{}".format(i),
                      [self._hparams.num_blocks, block_v_size],
                      initializer=tf.constant_initializer(0),
                      trainable=False,
                  )
                  ema_count.append(ema_count_i)
              with tf.colocate_with(means):
                  ema_means = []
                  for i in range(self._hparams.num_residuals):
                      ema_means_i = tf.get_variable(
                          "ema_means_{}".format(i),
                          [
                              self._hparams.num_blocks,
                              block_v_size,
                              block_dim,
                          ],
                          initializer=(
                              lambda shape, dtype=None, partition_info=None, verify_shape=None: means.initialized_value()[  # pylint: disable=g-long-lambda
                                  i
                              ]
                          ),
                          trainable=False,
                      )
                      ema_means.append(ema_means_i)
  
          # Update bottleneck
          self._hparams.bottleneck = functools.partial(
              self._hparams.bottleneck,
              means=means,
              ema_count=ema_count,
              ema_means=ema_means,
          )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import math
  import os
  from six.moves import range
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import common_image_attention
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import discretization
  from tensor2tensor.layers import latent_layers
  from tensor2tensor.layers import modalities
  from tensor2tensor.models import transformer
  from tensor2tensor.utils import beam_search
  from tensor2tensor.utils import contrib
  from tensor2tensor.utils import expert_utils
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import t2t_model
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 1 to create a new TransformerAE object from inherited base classes: t2t_model with example: obj = TransformerAE()