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.

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.,
            "extra_loss": 0.,
            "latent_pred": 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.,
              "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