class_docstr: Attention net.  See file docstring.
class_docstr_tok: ["Attention","net",".","See","file","docstring","."]
code: def model_fn(self, features):
      with tf.variable_scope(
          tf.get_variable_scope(), use_resource=True
      ) as vs:
          self._add_variable_scope("model_fn", vs)
          transformed_features = self.bottom(features)
  
          if self.hparams.activation_dtype == "bfloat16":
              for k, v in sorted(six.iteritems(transformed_features)):
                  if v.dtype == tf.float32:
                      transformed_features[k] = tf.cast(v, tf.bfloat16)
  
          with tf.variable_scope("body") as body_vs:
              self._add_variable_scope("body", body_vs)
              body_out = self.body(transformed_features)
          output, losses = self._normalize_body_output(body_out)
  
          if "training" in losses:
              tf.logging.info(
                  "Skipping T2TModel top and loss because training loss returned from body"
              )
              logits = output
          else:
              tf.logging.warn(
                  "The loss will be computed in model_fn now."
              )
              logits = self.top(output, features)
              losses["training"] = 0.0
              cur_kb_loss = losses["kb_loss"]
              cur_knowledge_training_loss = losses["transe_loss"]
              cur_kb_loss_weight = self._hparams.kb_loss_weight
              kb_train_weight = self._hparams.kb_train_weight
              cur_lm_loss_weight = 1.0 - cur_kb_loss_weight
              # Finalize loss
              if (
                  self._hparams.mode != tf_estimator.ModeKeys.PREDICT
                  and self._hparams.mode != "attack"
              ):
                  lm_loss_num, lm_loss_denom = self.loss(
                      logits, features
                  )
                  total_loss = (
                      kb_train_weight
                  ) * cur_knowledge_training_loss + (
                      1 - kb_train_weight
                  ) * (
                      cur_kb_loss * cur_kb_loss_weight
                      + (lm_loss_num / lm_loss_denom)
                      * cur_lm_loss_weight
                  )
                  tf.summary.scalar("kb_loss", cur_kb_loss)
                  tf.summary.scalar(
                      "transe_loss", cur_knowledge_training_loss
                  )
                  tf.summary.scalar(
                      "lm_loss", (lm_loss_num / lm_loss_denom)
                  )
                  tf.summary.scalar(
                      "cur_kb_loss_weight",
                      tf.reshape(cur_kb_loss_weight, []),
                  )
                  tf.logging.info("Loss computed " + str(total_loss))
                  losses = {"training": total_loss}
  
          return logits, losses
  
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 2 for calling the NeuralAssistant obj's underlying member method code with example usage: obj.model_fn(features) and returns: logits,  losses