class_docstr: Autoencoder-augmented Transformer.
class_docstr_tok: ["Autoencoder-augmented","Transformer","."]
docstr: Produce predictions from the model.
docstr_tok: ["Produce","predictions","from","the","model","."]
code: def infer(
      self,
      features=None,
      decode_length=50,
      beam_size=1,
      top_beams=1,
      alpha=0.0,
      use_tpu=False,
  ):
      """Produce predictions from the model."""
      if not self._hparams.do_mask:
          infer_out = super(TransformerAE, self).infer(
              features,
              decode_length,
              beam_size,
              top_beams,
              alpha,
              use_tpu=use_tpu,
          )
          return infer_out["outputs"]
      if not features:
          features = {}
      inputs_old = None
      if "inputs" in features and len(features["inputs"].shape) < 4:
          inputs_old = features["inputs"]
          features["inputs"] = tf.expand_dims(features["inputs"], 2)
  
      # Create an initial targets tensor.
      if "partial_targets" in features:
          initial_output = tf.convert_to_tensor(
              features["partial_targets"]
          )
      else:
          # inputs might not be present in features (e.g.: language modeling),
          # in which case we fallback to 'infer_targets' for calculating initial
          # input shape, type, etc.
          inputs_or_targets = features.get(
              "inputs", features.get("infer_targets")
          )
          batch_size = common_layers.shape_list(inputs_or_targets)[0]
          length = common_layers.shape_list(inputs_or_targets)[1]
          hidden_dim = common_layers.shape_list(inputs_or_targets)[-1]
          target_length = tf.to_int32(2.0 * tf.to_float(length))
          initial_output = tf.zeros(
              (batch_size, target_length, 1, hidden_dim),
              dtype=inputs_or_targets.dtype,
          )
  
      features["targets"] = initial_output
      logits, _ = self(features)  # pylint: disable=not-callable
      # this should only happen if we're doing target_modality not real
      if inputs_or_targets.dtype == tf.float32:
          samples = logits
      else:
          samples = tf.argmax(logits, axis=-1)
  
      # More steps.
      self.predict_mask = 0.0  # Use the provided targets this time.
      how_many_more_steps = (
          0  # Set to 1 or more for Gibbs-like sampling.
      )
      for _ in range(how_many_more_steps):
          with tf.variable_scope(tf.get_variable_scope(), reuse=True):
              features["targets"] = samples
              logits, _ = self(features)  # pylint: disable=not-callable
              if inputs_or_targets.dtype == tf.float32:
                  # When target_modality is real, the last axis does not represent
                  # classes, so it should not be argmax'ed
                  samples = logits
              else:
                  samples = tf.argmax(logits, axis=-1)
  
      self.predict_mask = 1.0
      if inputs_old is not None:  # Restore to not confuse Estimator.
          features["inputs"] = inputs_old
      return samples
  
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 4 for calling the TransformerAE obj's underlying member method code with example usage: obj.infer(features, decode_length, beam_size, top_beams, alpha, use_tpu) and returns: samples, infer_out