class_docstr: Universal Transformer: Depth-wise recurrent transformer model.
class_docstr_tok: ["Universal","Transformer",":","Depth-wise","recurrent","transformer","model","."]
docstr: Universal Transformer main model_fn.
  
  
  Args:
    features: Map of features to the model. Should contain the following:
        "inputs": Transformer inputs [batch_size, input_length, hidden_dim]
        "targets": Target decoder outputs.
            [batch_size, decoder_length, hidden_dim]
        "target_space_id"
  
  Returns:
    Final decoder representation. [batch_size, decoder_length, hidden_dim]
docstr_tok: ["Universal","Transformer","main","model_fn",".","Args",":","features",":","Map","of","features","to","the","model",".","Should","contain","the","following",":","``","inputs","''",":","Transformer","inputs","[","batch_size",",","input_length",",","hidden_dim","]","``","targets","''",":","Target","decoder","outputs",".","[","batch_size",",","decoder_length",",","hidden_dim","]","``","target_space_id","''","Returns",":","Final","decoder","representation",".","[","batch_size",",","decoder_length",",","hidden_dim","]"]
code: def body(self, features):
      """Universal Transformer main model_fn.
  
  
      Args:
        features: Map of features to the model. Should contain the following:
            "inputs": Transformer inputs [batch_size, input_length, hidden_dim]
            "targets": Target decoder outputs.
                [batch_size, decoder_length, hidden_dim]
            "target_space_id"
  
      Returns:
        Final decoder representation. [batch_size, decoder_length, hidden_dim]
      """
      hparams = self._hparams
      if hparams.add_position_timing_signal:
          # Turning off addition of positional embedding in the encoder/decoder
          # preparation as we do it in the beginning of each step.
          hparams.pos = None
  
      if self.has_input:
          inputs = features["inputs"]
          target_space = features["target_space_id"]
          (
              encoder_output,
              encoder_decoder_attention_bias,
              enc_extra_output,
          ) = self.encode(
              inputs, target_space, hparams, features=features
          )
      else:
          (
              encoder_output,
              encoder_decoder_attention_bias,
              enc_extra_output,
          ) = (None, None, (None, None))
  
      targets = features["targets"]
      targets = common_layers.flatten4d3d(targets)
  
      (
          decoder_input,
          decoder_self_attention_bias,
      ) = transformer.transformer_prepare_decoder(
          targets, hparams, features=features
      )
  
      decoder_output, dec_extra_output = self.decode(
          decoder_input,
          encoder_output,
          encoder_decoder_attention_bias,
          decoder_self_attention_bias,
          hparams,
          nonpadding=transformer.features_to_nonpadding(
              features, "targets"
          ),
      )
  
      expected_attentions = features.get("expected_attentions")
      if expected_attentions is not None:
          attention_loss = (
              common_attention.encoder_decoder_attention_loss(
                  expected_attentions,
                  self.attention_weights,
                  hparams.expected_attention_loss_type,
                  hparams.expected_attention_loss_multiplier,
              )
          )
          return decoder_output, {"attention_loss": attention_loss}
  
      if (
          hparams.recurrence_type == "act"
          and hparams.act_loss_weight != 0
      ):
          if self.has_input:
              enc_ponder_times, enc_remainders = enc_extra_output
              enc_act_loss = hparams.act_loss_weight * tf.reduce_mean(
                  enc_ponder_times + enc_remainders
              )
          else:
              enc_act_loss = 0.0
  
          (dec_ponder_times, dec_remainders) = dec_extra_output
          dec_act_loss = hparams.act_loss_weight * tf.reduce_mean(
              dec_ponder_times + dec_remainders
          )
          act_loss = enc_act_loss + dec_act_loss
          contrib.summary().scalar("act_loss", act_loss)
          return decoder_output, {"act_loss": act_loss}
  
      return decoder_output
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import common_layers
  from tensor2tensor.models import transformer
  from tensor2tensor.models.research import universal_transformer_util
  from tensor2tensor.utils import contrib
  from tensor2tensor.utils import registry
  import tensorflow.compat.v1
label_desc: Use image node_id 3 for calling the UniversalTransformer obj's underlying member method code with example usage: obj.body(features) and returns: decoder_output, decoder_output,  dict, decoder_output,  dict