class_docstr: Attention net. See file docstring.
class_docstr_tok: ["Attention","net",".","See","file","docstring","."]
docstr: Transformer main model_fn.
  
  Args:
    features: Map of features to the model. Should contain the following:
        "inputs": Transformer inputs.
            [batch_size, input_length, 1, hidden_dim].
        "targets": Target decoder outputs.
            [batch_size, target_length, 1, hidden_dim]
        "target_space_id": A scalar int from data_generators.problem.SpaceID.
  
  Returns:
    A 2-tuple containing:
        Logit tensor. [batch_size, decoder_length, vocab_size]
        Map of keys to loss tensors. Should contain the following:
            "training": Training loss (shift == 0).
            "auxiliary": Auxiliary loss (shift != 0).
docstr_tok: ["Transformer","main","model_fn",".","Args",":","features",":","Map","of","features","to","the","model",".","Should","contain","the","following",":","``","inputs","''",":","Transformer","inputs",".","[","batch_size",",","input_length",",","1",",","hidden_dim","]",".","``","targets","''",":","Target","decoder","outputs",".","[","batch_size",",","target_length",",","1",",","hidden_dim","]","``","target_space_id","''",":","A","scalar","int","from","data_generators.problem.SpaceID",".","Returns",":","A","2-tuple","containing",":","Logit","tensor",".","[","batch_size",",","decoder_length",",","vocab_size","]","Map","of","keys","to","loss","tensors",".","Should","contain","the","following",":","``","training","''",":","Training","loss","(","shift","==","0",")",".","``","auxiliary","''",":","Auxiliary","loss","(","shift","!","=","0",")","."]
code: def body(self, features):
      """Transformer main model_fn.
  
      Args:
        features: Map of features to the model. Should contain the following:
            "inputs": Transformer inputs.
                [batch_size, input_length, 1, hidden_dim].
            "targets": Target decoder outputs.
                [batch_size, target_length, 1, hidden_dim]
            "target_space_id": A scalar int from data_generators.problem.SpaceID.
  
      Returns:
        A 2-tuple containing:
            Logit tensor. [batch_size, decoder_length, vocab_size]
            Map of keys to loss tensors. Should contain the following:
                "training": Training loss (shift == 0).
                "auxiliary": Auxiliary loss (shift != 0).
      """
      output = super(TransformerAux, self).body(features)
      output, losses = self._normalize_body_output(output)
  
      aux = 0.0
      for shift in self._extract_shift_values():
          loss_num, loss_den = self.auxiliary_loss(
              output, features, shift
          )
          aux += loss_num / loss_den
      losses["auxiliary"] = aux
  
      return output, losses
  
file_path: tensor2tensor/tensor2tensor/models/research/transformer_aux.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from tensor2tensor.layers import common_layers
  from tensor2tensor.models import transformer
  from tensor2tensor.utils import registry
  import tensorflow.compat.v1
label_desc: Use image node_id 3 for calling the TransformerAux obj's underlying member method code with example usage: obj.body(features) and returns: output,  losses