code: def mnist_model(image, labels, mesh):
      """The model.
  
      Args:
        image: tf.Tensor with shape [batch, 28*28]
        labels: a tf.Tensor with shape [batch] and dtype tf.int32
        mesh: a mtf.Mesh
  
      Returns:
        logits: a mtf.Tensor with shape [batch, 10]
        loss: a mtf.Tensor with shape []
      """
      batch_dim = mtf.Dimension("batch", FLAGS.batch_size)
      row_blocks_dim = mtf.Dimension("row_blocks", 4)
      col_blocks_dim = mtf.Dimension("col_blocks", 4)
      rows_dim = mtf.Dimension("rows_size", 7)
      cols_dim = mtf.Dimension("cols_size", 7)
  
      classes_dim = mtf.Dimension("classes", 10)
      one_channel_dim = mtf.Dimension("one_channel", 1)
  
      x = mtf.import_tf_tensor(
          mesh,
          tf.reshape(image, [FLAGS.batch_size, 4, 7, 4, 7, 1]),
          mtf.Shape(
              [
                  batch_dim,
                  row_blocks_dim,
                  rows_dim,
                  col_blocks_dim,
                  cols_dim,
                  one_channel_dim,
              ]
          ),
      )
      x = mtf.transpose(
          x,
          [
              batch_dim,
              row_blocks_dim,
              col_blocks_dim,
              rows_dim,
              cols_dim,
              one_channel_dim,
          ],
      )
  
      # add some convolutional layers to demonstrate that convolution works.
      filters1_dim = mtf.Dimension("filters1", 16)
      filters2_dim = mtf.Dimension("filters2", 16)
      f1 = mtf.relu(
          mtf.layers.conv2d_with_blocks(
              x,
              filters1_dim,
              filter_size=[9, 9],
              strides=[1, 1],
              padding="SAME",
              h_blocks_dim=row_blocks_dim,
              w_blocks_dim=col_blocks_dim,
              name="conv0",
          )
      )
      f2 = mtf.relu(
          mtf.layers.conv2d_with_blocks(
              f1,
              filters2_dim,
              filter_size=[9, 9],
              strides=[1, 1],
              padding="SAME",
              h_blocks_dim=row_blocks_dim,
              w_blocks_dim=col_blocks_dim,
              name="conv1",
          )
      )
      x = mtf.reduce_mean(f2, reduced_dim=filters2_dim)
  
      # add some fully-connected dense layers.
      hidden_dim1 = mtf.Dimension("hidden1", FLAGS.hidden_size)
      hidden_dim2 = mtf.Dimension("hidden2", FLAGS.hidden_size)
  
      h1 = mtf.layers.dense(
          x,
          hidden_dim1,
          reduced_dims=x.shape.dims[-4:],
          activation=mtf.relu,
          name="hidden1",
      )
      h2 = mtf.layers.dense(
          h1, hidden_dim2, activation=mtf.relu, name="hidden2"
      )
      logits = mtf.layers.dense(h2, classes_dim, name="logits")
      if labels is None:
          loss = None
      else:
          labels = mtf.import_tf_tensor(
              mesh,
              tf.reshape(labels, [FLAGS.batch_size]),
              mtf.Shape([batch_dim]),
          )
          loss = mtf.layers.softmax_cross_entropy_with_logits(
              logits, mtf.one_hot(labels, classes_dim), classes_dim
          )
          loss = mtf.reduce_mean(loss)
      return logits, loss
  
file_path: mesh/examples/mnist.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import mesh_tensorflow
  import mnist_dataset
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 1 for calling a global function with example usage: mnist_model(image, labels, mesh) and returns: logits,  loss