class_docstr: Image Transformer in mesh_tensorflow.
class_docstr_tok: ["Image","Transformer","in","mesh_tensorflow","."]
code: def mtf_model_fn(self, features, mesh):
      features = copy.copy(features)
      tf.logging.info("features = %s" % features)
      hparams = self._hparams
      activation_dtype = self.activation_type
  
      # We assume fixed vocab size for targets
      targets = tf.to_int32(features["targets"])
  
      # Image preprocessing, reshape into a 1D sequence and shift right.
      length = hparams.img_len * hparams.img_len * hparams.num_channels
      targets = tf.reshape(targets, [hparams.batch_size, length])
      shifted_targets = common_layers.shift_right_2d(targets)
  
      # Declare all the dimensions
      batch_dim = mtf.Dimension("batch", hparams.batch_size)
  
      def import_to_batch_by_length(x, name):
          return mtf.import_tf_tensor(
              mesh,
              x,
              mtf.Shape([batch_dim, self.length_dim]),
              name=name,
          )
  
      targets = import_to_batch_by_length(targets, "targets")
      shifted_targets = import_to_batch_by_length(
          shifted_targets, "shifted_targets"
      )
  
      extra_losses = []
  
      # Create targets content and position embeddings.
      # Create embedding var for targets and positions and do a gather.
      targets_embedding_var = mtf.get_variable(
          mesh,
          "targets_embedding",
          mtf.Shape([self.targets_vocab_dim, self.model_dim]),
          initializer=tf.random_normal_initializer(),
          activation_dtype=activation_dtype,
      )
  
      x = mtf.gather(
          targets_embedding_var, shifted_targets, self.targets_vocab_dim
      )
  
      # Add positional embeddings
      x += mtf.reshape(
          self.create_positional_emb_2d(targets),
          [self.length_dim, self.model_dim],
      )
  
      # If conditional and input is given, add the input embedding to the target.
      # TODO(nikip): Verify conditional.
      if self.has_input and not hparams.unconditional:
          inputs = tf.squeeze(tf.to_int32(features["inputs"]), [2, 3])
          inputs = import_to_batch_by_length(inputs, "inputs")
  
          # Input embeddings
          inputs_embedding_var = mtf.layers.embedding(
              mesh,
              "input_embedding",
              mtf.Shape([self.inputs_vocab_dim, self.model_dim]),
              activation_dtype=activation_dtype,
          )
          inputs_emb = mtf.gather(
              inputs_embedding_var, inputs, self.inputs_vocab_dim
          )
          x += inputs_emb
  
      # Image Transformer Decoder
      # [ self attention - ffn - residual + dropout] x n
      if hparams.attention_type == "local1d_spatial":
          decoder_output = local_attention1d_spatial_decoder(
              x,
              self.kv_dim,
              self.heads_dim,
              self.feedforward_dim,
              hparams,
          )
      elif hparams.attention_type == "local2d_spatial":
          decoder_output = local_attention2d_spatial_decoder(
              x,
              self.kv_dim,
              self.heads_dim,
              self.feedforward_dim,
              hparams,
          )
      elif hparams.attention_type == "local1d":
          decoder_output = local_attention1d_masked_decoder(
              x,
              self.kv_dim,
              self.heads_dim,
              self.feedforward_dim,
              hparams,
          )
      else:
          raise ValueError("Invalid attention type.")
  
      # Calculate the logits and loss.
      logits = mtf.layers.dense(
          decoder_output, self.outputs_vocab_dim, name="logits"
      )
      # Need a reshape for logits
      logits = mtf.reshape(
          logits,
          mtf.Shape(
              [batch_dim, self.length_dim, self.outputs_vocab_dim]
          ),
      )
      soft_targets = mtf.one_hot(
          targets, self.outputs_vocab_dim, dtype=activation_dtype
      )
      loss = mtf.layers.softmax_cross_entropy_with_logits(
          logits, soft_targets, self.outputs_vocab_dim
      )
      loss = mtf.reduce_mean(loss)
      for l in extra_losses:
          loss += l
  
      # Reshape logits to original target shape.
      logits = mtf.reshape(
          logits,
          mtf.Shape(
              [
                  batch_dim,
                  self.rows_dim,
                  self.orig_cols_dim,
                  self.channels_dim,
                  self.outputs_vocab_dim,
              ]
          ),
      )
  
      return logits, loss
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  import mesh_tensorflow
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.layers import common_layers
  from tensor2tensor.utils import mtf_model
  from tensor2tensor.utils import registry
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 17 for calling the MtfImageTransformer obj's underlying member method code with example usage: obj.mtf_model_fn(features, mesh) and returns: logits,  loss, mtf