code: def model_fn(features, labels, mode, params):
      """A model is called by TpuEstimator."""
      del labels
      global_step = tf.train.get_global_step()
      graph = mtf.Graph()
      mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
      layout_rules = mtf.convert_to_layout_rules(FLAGS.layout)
      if FLAGS.use_tpu:
          ctx = params["context"]
          num_hosts = ctx.num_hosts
          host_placement_fn = ctx.tpu_host_placement_function
          device_list = [
              host_placement_fn(host_id=t) for t in range(num_hosts)
          ]
          tf.logging.info(
              "device_list = %s" % device_list,
          )
          # TODO(ylc): Better estimation of replica cache size?
          replica_cache_size = 300 * 1000000  # 300M per replica
          # Worker 0 caches all the TPU binaries.
          worker0_mem = replica_cache_size * ctx.num_replicas
          devices_memeory_usage = [worker0_mem] + [0] * (num_hosts - 1)
          var_placer = mtf.utils.BalancedVariablePlacer(
              device_list, devices_memeory_usage
          )
          mesh_devices = [""] * mesh_shape.size
          mesh_impl = mtf.simd_mesh_impl.SimdMeshImpl(
              mesh_shape,
              layout_rules,
              mesh_devices,
              ctx.device_assignment,
          )
      else:
          var_placer = None
          mesh_devices = [""] * mesh_shape.size
          mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
              mesh_shape, layout_rules, mesh_devices
          )
      mesh = mtf.Mesh(graph, "my_mesh", var_placer)
  
      with mtf.utils.outside_all_rewrites():
          logits, loss = toy_model(features, mesh)
  
      # TRAIN mode
      if mode == tf_estimator.ModeKeys.TRAIN:
          var_grads = mtf.gradients(
              [loss], [v.outputs[0] for v in graph.trainable_variables]
          )
          if FLAGS.optimizer == "Adafactor":
              optimizer = mtf.optimize.AdafactorOptimizer()
          else:
              assert FLAGS.optimizer == "SGD"
              optimizer = mtf.optimize.SgdOptimizer(
                  learning_rate=FLAGS.lr
              )
          update_ops = optimizer.apply_grads(
              var_grads, graph.trainable_variables
          )
      else:
          # for now, we can only export fully-replicated tensors.
          fully_replicated_logits = mtf.anonymize(logits)
  
      lowering = mtf.Lowering(graph, {mesh: mesh_impl})
  
      tf_loss = tf.to_float(lowering.export_to_tf_tensor(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))
          tf.logging.info("tf_update_ops: {}".format(tf_update_ops))
          train_op = tf.group(tf_update_ops)
      else:
          tf_logits = lowering.export_to_tf_tensor(
              fully_replicated_logits
          )
  
      with mtf.utils.outside_all_rewrites():
          # Copy master variables to slices. Must be called first.
          restore_hook = mtf.MtfRestoreHook(lowering)
          if mode == tf_estimator.ModeKeys.TRAIN:
              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],
              )
  
              return tpu_estimator.TPUEstimatorSpec(
                  tf_estimator.ModeKeys.TRAIN,
                  loss=tf_loss,
                  train_op=train_op,
                  training_hooks=[restore_hook, saver_hook],
              )
          elif mode == tf_estimator.ModeKeys.EVAL:
  
              def metric_fn(tf_logits):
                  mean_logits = tf.metrics.mean(tf_logits)
                  return {"mean_logits": mean_logits}
  
              eval_metrics = (metric_fn, [tf_logits])
  
              return tpu_estimator.TPUEstimatorSpec(
                  tf_estimator.ModeKeys.EVAL,
                  evaluation_hooks=[restore_hook],
                  loss=tf_loss,
                  eval_metrics=eval_metrics,
              )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import mesh_tensorflow
  import numpy
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  from tensorflow.python.data.ops.dataset_ops import Dataset
  from tensorflow.python.platform import flags
  from tensorflow.python.platform import tf_logging
  from tensorflow_estimator.python.estimator import estimator
  from tensorflow_estimator.python.estimator.tpu import tpu_config
  from tensorflow_estimator.python.estimator.tpu import tpu_estimator
label_desc: Use image node_id 2 for calling a global function with example usage: model_fn(features, labels, mode, params) and returns: tpu_estimator, tpu_estimator, dict