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, hidden_dim]
        "targets": Target decoder outputs. [batch_size, decoder_length,
          hidden_dim]
        "target_space_id": A scalar int from data_generators.problem.SpaceID.
  
  Returns:
    Final decoder representation. [batch_size, decoder_length, hidden_dim]
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",",","hidden_dim","]","``","targets","''",":","Target","decoder","outputs",".","[","batch_size",",","decoder_length",",","hidden_dim","]","``","target_space_id","''",":","A","scalar","int","from","data_generators.problem.SpaceID",".","Returns",":","Final","decoder","representation",".","[","batch_size",",","decoder_length",",","hidden_dim","]"]
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, hidden_dim]
            "targets": Target decoder outputs. [batch_size, decoder_length,
              hidden_dim]
            "target_space_id": A scalar int from data_generators.problem.SpaceID.
  
      Returns:
        Final decoder representation. [batch_size, decoder_length, hidden_dim]
      """
      tf.logging.info("Using PgScratch BODY function.")
      hparams = self._hparams
  
      losses = {}
      inputs = features["inputs"]
      target_space = features["target_space_id"]
      # encoder_output: <tf.float32>[batch_size, input_length, hidden_dim]
      # encoder_decoder_attention_bias: <tf.float32>[batch_size, input_length]
      encoder_output, encoder_decoder_attention_bias = self.encode(
          inputs,
          target_space,
          hparams,
          features=features,
          losses=losses,
      )
  
      with tf.variable_scope("knowledge"):
          with tf.name_scope("knowledge_encoding"):
              # Encode knowledge.
              # <tf.float32>[batch_size, triple_num, emb_dim]
              (
                  fact_embedding,
                  fact_lengths,
              ) = self.encode_knowledge_bottom(features)
              tf.logging.info("Encoded knowledge")
  
          with tf.name_scope("knowledge_selection_and_loss"):
              # Compute knowledge selection and loss.
              (
                  triple_logits,
                  avg_triple_selection_loss,
                  knowledge_encoder_output,
                  transe_loss,
              ) = self.compute_knowledge_selection_and_loss(
                  features,
                  encoder_output,
                  fact_embedding,
                  fact_lengths,
                  hparams.margin,
                  hparams.num_negative_samples,
              )
              losses["kb_loss"] = avg_triple_selection_loss
              losses["transe_loss"] = transe_loss
  
      if hparams.attend_kb:
          tf.logging.info("ATTEND_KB is ACTIVE")
          with tf.name_scope("knowledge_attention"):
              knowledge_padding = tf.zeros_like(
                  triple_logits, dtype=tf.float32
              )
              knowledge_attention_bias = (
                  common_attention.attention_bias_ignore_padding(
                      knowledge_padding
                  )
              )
              encoder_output = tf.concat(
                  [knowledge_encoder_output, encoder_output], 1
              )
              encoder_decoder_attention_bias = tf.concat(
                  [
                      knowledge_attention_bias,
                      encoder_decoder_attention_bias,
                  ],
                  -1,
              )
  
      else:
          tf.logging.info("ATTEND_KB is INACTIVE")
  
      targets = features["targets"]
      targets_shape = common_layers.shape_list(targets)
      targets = common_layers.flatten4d3d(targets)
  
      (
          decoder_input,
          decoder_self_attention_bias,
      ) = transformer.transformer_prepare_decoder(
          targets, hparams, features=features
      )
  
      decode_kwargs = {}
      decoder_output = self.decode(
          decoder_input,
          encoder_output,
          encoder_decoder_attention_bias,
          decoder_self_attention_bias,
          hparams,
          nonpadding=transformer.features_to_nonpadding(
              features, "targets"
          ),
          losses=losses,
          **decode_kwargs,
      )
  
      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}
  
      ret = tf.reshape(decoder_output, targets_shape)
      if losses:
          return ret, losses
      else:
          return ret
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import six
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import common_layers
  from tensor2tensor.models import transformer
  from tensor2tensor.utils import registry
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 5 for calling the NeuralAssistant obj's underlying member method code with example usage: obj.body(features) and returns: decoder_output,  dict, ret,  losses, ret