class_docstr: ResNet 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.set_activation_type()
      is_training = hparams.mode == tf_estimator.ModeKeys.TRAIN
  
      # Declare all the dimensions
      batch_dim = mtf.Dimension("batch", hparams.batch_size)
      hidden_dim = mtf.Dimension("hidden", hparams.hidden_size)
      filter_dim = mtf.Dimension("filters", hparams.filter_sizes[0])
      rows_dim = mtf.Dimension("rows_size", hparams.rows_size)
      cols_dim = mtf.Dimension("cols_size", hparams.cols_size)
      row_blocks_dim = mtf.Dimension("row_blocks", hparams.row_blocks)
      col_blocks_dim = mtf.Dimension("col_blocks", hparams.col_blocks)
      classes_dim = mtf.Dimension("classes", 10)
      channels_dim = mtf.Dimension("channels", 3)
      one_channel_dim = mtf.Dimension("one_channel", 1)
  
      inputs = features["inputs"]
      x = mtf.import_tf_tensor(
          mesh,
          tf.reshape(
              inputs,
              [
                  hparams.batch_size,
                  hparams.row_blocks,
                  hparams.rows_size // hparams.row_blocks,
                  hparams.col_blocks,
                  hparams.num_channels
                  * hparams.cols_size
                  // hparams.col_blocks,
                  hparams.num_channels,
              ],
          ),
          mtf.Shape(
              [
                  batch_dim,
                  row_blocks_dim,
                  rows_dim,
                  col_blocks_dim,
                  cols_dim,
                  channels_dim,
              ]
          ),
      )
      x = mtf.transpose(
          x,
          [
              batch_dim,
              row_blocks_dim,
              col_blocks_dim,
              rows_dim,
              cols_dim,
              channels_dim,
          ],
      )
  
      x = mtf.to_float(x)
      x = mtf.layers.conv2d_with_blocks(
          x,
          filter_dim,
          filter_size=[3, 3],
          strides=[1, 1],
          padding="SAME",
          h_blocks_dim=None,
          w_blocks_dim=col_blocks_dim,
          name="initial_filter",
      )
  
      x = batch_norm_relu(x, is_training)
  
      # Conv blocks
      # [block - strided block layer - strided block layer] x n
      for layer in range(hparams.num_layers):
          layer_name = "block_layer_%d" % layer
          with tf.variable_scope(layer_name):
              # Residual block layer
              x = block_layer(
                  inputs=x,
                  filters=hparams.filter_sizes[0],
                  blocks=hparams.layer_sizes[0],
                  strides=[1, 1],
                  is_training=is_training,
                  name="block_layer1",
                  row_blocks_dim=None,
                  col_blocks_dim=None,
              )
              x = block_layer(
                  inputs=x,
                  filters=hparams.filter_sizes[1],
                  blocks=hparams.layer_sizes[1],
                  strides=[1, 1],
                  is_training=is_training,
                  name="block_layer2",
                  row_blocks_dim=None,
                  col_blocks_dim=None,
              )
              x = block_layer(
                  inputs=x,
                  filters=hparams.filter_sizes[2],
                  blocks=hparams.layer_sizes[2],
                  strides=[1, 1],
                  is_training=is_training,
                  name="block_layer3",
                  row_blocks_dim=None,
                  col_blocks_dim=None,
              )
  
      # Calculate the logits and loss.
      out = x
      outputs = mtf.layers.dense(
          out,
          hidden_dim,
          reduced_dims=out.shape.dims[-5:],
          activation=mtf.relu,
          name="dense",
      )
  
      # We assume fixed vocab size for targets
      labels = tf.squeeze(tf.to_int32(features["targets"]), [2, 3])
      labels = mtf.import_tf_tensor(
          mesh,
          tf.reshape(labels, [hparams.batch_size]),
          mtf.Shape([batch_dim]),
      )
  
      logits = mtf.layers.dense(outputs, classes_dim, name="logits")
      soft_targets = mtf.one_hot(
          labels, classes_dim, dtype=activation_dtype
      )
      loss = mtf.layers.softmax_cross_entropy_with_logits(
          logits, soft_targets, classes_dim
      )
  
      # Reshape logits so it doesn't break inside t2t.
      logits = mtf.reshape(
          logits, mtf.Shape([batch_dim, one_channel_dim, classes_dim])
      )
      loss = mtf.reduce_mean(loss)
      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.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 2 for calling the MtfResNet obj's underlying member method code with example usage: obj.mtf_model_fn(features, mesh) and returns: logits,  loss