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
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]
    
tensor2tensor/models/neural_assistant.py