code: def transformer_autoencoder(
      inputs,
      targets,
      target_space,
      hparams,
      cache=None,
      predict_mask=1.0,
  ):
      """Auto-encoder using a Transformer decoder and a prior over latent sequences.
  
      Args:
        inputs: Tensor of shape [batch, length, 1, hparams.hidden_size] or None.
        targets: Tensor of shape [batch, ..., channels]. Ellipses may be 1 or 2
          dimensions denoting sequence length.
        target_space: int. Used for encoding inputs under a target space id.
        hparams: HParams.
        cache: Tensor of shape [batch, length] or None.
        predict_mask: Tensor masking whether to use gold targets or predictions.
  
      Returns:
        decoder_output: Tensor of shape [batch, ..., hparams.hidden_size] presenting
          pre-logit activations. After a transformation (`top` in `T2TModel`), it is
          used with targets to compute the "training" (reconstruction) loss.
        losses: dict of str to Tensors. There are three loss terms: "extra",
          "extra_loss", and "latent_pred". The first is hard-coded to 0. The latter
          two are Tensors of shape [batch].
        cache: Tensor of shape [batch, length], either the same as cache, or newly
          computed if the cache input is None.
      """
      original_targets_shape = common_layers.shape_list(targets)
      batch_size = original_targets_shape[0]
      if len(original_targets_shape) == 4:
          compress_fn = compress_encoder_2d
          decompress_fn = decompress_decoder_2d
      else:
          compress_fn = compress_encoder_1d
          decompress_fn = decompress_decoder_1d
  
      ed_attention_bias = None
      if inputs is not None:
          inputs, ed_attention_bias = transformer_text_encoder(
              inputs, target_space, hparams, name="input_encoder"
          )
  
      losses = {"extra": 0.0, "extra_loss": 0.0, "latent_pred": 0.0}
      if hparams.mode != tf_estimator.ModeKeys.PREDICT:
          targets_compressed = compress_fn(
              targets, hparams, name="compress"
          )
  
          if hparams.mode == tf_estimator.ModeKeys.TRAIN:
              scale = common_layers.inverse_exp_decay(
                  hparams.startup_steps
              )
          else:
              scale = 1.0
          scale = tf.to_float(
              tf.less(tf.random_uniform([batch_size]), scale)
          )
  
          (
              latents_dense,
              latents_discrete,
              extra_loss,
              _,
          ) = bottleneck_layer(targets_compressed, hparams)
          extra_loss = scale * tf.reduce_mean(extra_loss)
  
          _, latents_pred_loss = latent_prediction_model(
              inputs,
              ed_attention_bias,
              latents_discrete,
              latents_dense,
              hparams,
              name="latent_pred",
          )
          latent_time = tf.less(
              hparams.mask_startup_steps,
              tf.to_int32(tf.train.get_global_step()),
          )
          latents_pred_loss = scale * tf.reduce_mean(latents_pred_loss)
          latents_pred_loss *= tf.to_float(latent_time)
  
          # Apply dropout noise for each data point and time step.
          latents_dense_shape = common_layers.shape_list(latents_dense)
          latents_dense = tf.nn.dropout(
              latents_dense,
              keep_prob=1 - hparams.latent_dropout,
              noise_shape=[
                  latents_dense_shape[0],
                  latents_dense_shape[1],
                  1,
              ],
          )
  
          # TODO(trandustin): Can we combine extra and extra_loss?
          losses = {
              "extra": 0.0,
              "extra_loss": extra_loss,
              "latent_pred": latents_pred_loss,
          }
      else:
          # Set the latent length, which is num_latents times the number of latent
          # pixels. The number of latent pixels is determined by a compression factor
          # on the number of image pixels.
          latent_len = (
              hparams.img_len * hparams.img_len * hparams.num_latents
          ) / (2**hparams.num_compress_steps)
          _, _, _, embed_fn = bottleneck_layer(
              targets_compressed, hparams
          )
          latents_dense = tf.zeros(
              [batch_size, latent_len, 1, hparams.hidden_size]
          )
          if cache is None:
              cache = ae_latent_sample_beam(
                  latents_dense,
                  inputs,
                  ed_attention_bias,
                  embed_fn,
                  hparams,
              )
          cache_one_hot = tf.one_hot(
              cache, depth=2**hparams.bottleneck_bits
          )
          latents_dense = embed_fn(cache_one_hot, hparams.hidden_size)
  
      if len(original_targets_shape) == 4:
          compressed_img_len = hparams.img_len // 2 ** (
              hparams.num_compress_steps // 2
          )
          latents_dense = tf.reshape(
              latents_dense,
              [
                  batch_size,
                  compressed_img_len,
                  compressed_img_len,
                  hparams.num_latents * hparams.hidden_size,
              ],
          )
  
      latents_dense = decompress_fn(
          latents_dense, hparams, name="decompress"
      )
      latents_dense = tf.reshape(
          latents_dense,
          [-1, hparams.img_len, hparams.img_len, hparams.hidden_size],
      )
  
      if hparams.use_gold_targets:
          if hparams.mode == tf_estimator.ModeKeys.PREDICT:
              masking = predict_mask
          else:
              masking = common_layers.inverse_exp_decay(
                  hparams.mask_startup_steps
              )
          targets, _, _ = cia.maybe_reshape_4d_to_3d(targets)
          mask = tf.less(
              masking,
              tf.random_uniform(common_layers.shape_list(targets)[:-1]),
          )
          mask = tf.expand_dims(tf.to_float(mask), 2)
          latents_dense = mask * targets + (1.0 - mask) * latents_dense
  
      latents_dense = tf.reshape(latents_dense, original_targets_shape)
      if hparams.decode_autoregressive:
          decoder_output = transformer_image_decoder(
              latents_dense,
              inputs,
              ed_attention_bias,
              hparams,
              name="decoder",
          )
      else:
          decoder_output = latents_dense
      return decoder_output, losses, cache
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  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 transformer_layers
  from tensor2tensor.utils import beam_search
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tensorflow_probability
label_desc: Use image node_id 18 for calling a global function with example usage: transformer_autoencoder(inputs, targets, target_space, hparams, cache, predict_mask) and returns: decoder_output,  losses,  cache