code: def model_fn(features, labels, mode, params):
      """The model_fn argument for creating an Estimator."""
      tf.logging.info(
          "features = %s labels = %s mode = %s params=%s"
          % (features, labels, mode, params)
      )
      global_step = tf.train.get_global_step()
      graph = mtf.Graph()
      mesh = mtf.Mesh(graph, "my_mesh")
      logits, loss = mnist_model(features, labels, mesh)
      mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
      layout_rules = mtf.convert_to_layout_rules(FLAGS.layout)
      mesh_size = mesh_shape.size
      mesh_devices = [""] * mesh_size
      mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
          mesh_shape, layout_rules, mesh_devices
      )
  
      if mode == tf_estimator.ModeKeys.TRAIN:
          var_grads = mtf.gradients(
              [loss], [v.outputs[0] for v in graph.trainable_variables]
          )
          optimizer = mtf.optimize.AdafactorOptimizer()
          update_ops = optimizer.apply_grads(
              var_grads, graph.trainable_variables
          )
  
      lowering = mtf.Lowering(graph, {mesh: mesh_impl})
      restore_hook = mtf.MtfRestoreHook(lowering)
  
      tf_logits = lowering.export_to_tf_tensor(logits)
      if mode != tf_estimator.ModeKeys.PREDICT:
          tf_loss = lowering.export_to_tf_tensor(loss)
          tf.summary.scalar("loss", tf_loss)
  
      if mode == tf_estimator.ModeKeys.TRAIN:
          tf_update_ops = [
              lowering.lowered_operation(op) for op in update_ops
          ]
          tf_update_ops.append(tf.assign_add(global_step, 1))
          train_op = tf.group(tf_update_ops)
          saver = tf.train.Saver(
              tf.global_variables(),
              sharded=True,
              max_to_keep=10,
              keep_checkpoint_every_n_hours=2,
              defer_build=False,
              save_relative_paths=True,
          )
          tf.add_to_collection(tf.GraphKeys.SAVERS, saver)
          saver_listener = mtf.MtfCheckpointSaverListener(lowering)
          saver_hook = tf.train.CheckpointSaverHook(
              FLAGS.model_dir,
              save_steps=1000,
              saver=saver,
              listeners=[saver_listener],
          )
  
          accuracy = tf.metrics.accuracy(
              labels=labels, predictions=tf.argmax(tf_logits, axis=1)
          )
  
          # Name tensors to be logged with LoggingTensorHook.
          tf.identity(tf_loss, "cross_entropy")
          tf.identity(accuracy[1], name="train_accuracy")
  
          # Save accuracy scalar to Tensorboard output.
          tf.summary.scalar("train_accuracy", accuracy[1])
  
          # restore_hook must come before saver_hook
          return tf_estimator.EstimatorSpec(
              tf_estimator.ModeKeys.TRAIN,
              loss=tf_loss,
              train_op=train_op,
              training_chief_hooks=[restore_hook, saver_hook],
          )
  
      if mode == tf_estimator.ModeKeys.PREDICT:
          predictions = {
              "classes": tf.argmax(tf_logits, axis=1),
              "probabilities": tf.nn.softmax(tf_logits),
          }
          return tf_estimator.EstimatorSpec(
              mode=tf_estimator.ModeKeys.PREDICT,
              predictions=predictions,
              prediction_hooks=[restore_hook],
              export_outputs={
                  "classify": tf_estimator.export.PredictOutput(
                      predictions
                  )
              },
          )
      if mode == tf_estimator.ModeKeys.EVAL:
          return tf_estimator.EstimatorSpec(
              mode=tf_estimator.ModeKeys.EVAL,
              loss=tf_loss,
              evaluation_hooks=[restore_hook],
              eval_metric_ops={
                  "accuracy": tf.metrics.accuracy(
                      labels=labels,
                      predictions=tf.argmax(tf_logits, axis=1),
                  ),
              },
          )
  
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 2 for calling a global function with example usage: model_fn(features, labels, mode, params) and returns: tf_estimator, tf_estimator, tf_estimator