code: def _get_model_fn(train_or_eval, mesh_context):
      """Returns _model_fn."""
      captured_hooks = _CapturedObject()
      captured_output_dtypes_shapes = _CapturedObject()
      assert train_or_eval in ["train", "eval"]
  
      def _model_fn(input_fea, input_lab):
          """Creates a model, add summary, modes (train or eval), and hooks."""
  
          # input_fea and input_lab should be a list (laid_out_tensors).
          if not isinstance(input_fea, list):
              input_fea = [input_fea]
          if not isinstance(input_lab, list):
              input_lab = [input_lab]
  
          def _add_summary(
              lowering, train_or_eval, tf_loss, scalars, global_step
          ):
              """Add all summaries."""
              for k in scalars.keys():
                  if not isinstance(scalars[k], tf.Tensor):
                      scalars[k] = tf.cast(
                          lowering.export_to_tf_tensor(scalars[k]),
                          tf.float32,
                      )
  
              def _host_loss_summary(global_step, tf_loss, **scalars):
                  """Add summary.scalar in host side."""
                  gs = tf.cast(global_step, tf.int64)
                  sum_loss = contrib_summary.scalar(
                      "{}_loss".format(train_or_eval), tf_loss, step=gs
                  )
                  sum_ops = [sum_loss.op]
                  for description, tf_metric in six.iteritems(scalars):
                      sum_metric = contrib_summary.scalar(
                          "{}_{}".format(train_or_eval, description),
                          tf_metric,
                          step=gs,
                      )
                      sum_ops.append(sum_metric)
                  with tf.control_dependencies(sum_ops):
                      return tf.identity(tf_loss)
  
              if FLAGS.use_tpu:
                  # Cast the global step to tf.int32, since
                  # outside_compilation does not support tf.int64.
                  tf_loss = tpu.outside_compilation(
                      _host_loss_summary,
                      tf.cast(global_step, tf.int32),
                      tf_loss,
                      **scalars,
                  )
              else:
                  tf_loss = _host_loss_summary(
                      tf.cast(global_step, tf.int32), tf_loss, **scalars
                  )
  
              return tf_loss
  
          global_step = tf.train.get_or_create_global_step()
          (
              graph,
              mesh,
              mesh_impl,
          ) = mesh_context.create_graph_mesh_and_mesh_impl()
  
          with mtf.utils.outside_all_rewrites():
              # Do not tpu_rewrite this part. Inside this unet, If you use Tensorflow,
              # instead of Mesh-Tensorflor, it will cause host to tpu send/rec.
              (
                  preds,
                  loss,
                  scalars,
                  bn_update_ops,
              ) = unet.unet_with_spatial_partition(
                  mesh, mesh_impl, train_or_eval, input_fea, input_lab
              )
  
          if train_or_eval == "train":
              var_grads = mtf.gradients(
                  [loss],
                  [v.outputs[0] for v in graph.trainable_variables],
              )
  
              lr = FLAGS.lr * tf.pow(
                  FLAGS.lr_drop_rate,
                  tf.floor(
                      tf.cast(global_step, tf.float32)
                      / FLAGS.lr_drop_steps
                  ),
              )
              scalars["learning_rate"] = lr
  
              optimizer = mtf.optimize.AdafactorOptimizer(
                  learning_rate=lr
              )
              update_ops = optimizer.apply_grads(
                  var_grads, graph.trainable_variables
              )
  
              # This is where the actual tf graph got built.
              lowering = mtf.Lowering(graph, {mesh: mesh_impl})
  
              tf_update_ops = [
                  lowering.lowered_operation(op) for op in update_ops
              ]
              tf_update_ops.append(tf.assign_add(global_step, 1))
              tf_update_ops.extend(
                  [
                      lowering.lowered_operation(op)
                      for op in bn_update_ops
                  ]
              )
  
          else:  # train_or_eval == 'eval':
              preds = [mtf.anonymize(pred) for pred in preds]
  
              # This is where the actual tf graph got built.
              lowering = mtf.Lowering(graph, {mesh: mesh_impl})
  
              tf_preds = [
                  tf.cast(
                      lowering.export_to_tf_tensor(pred), tf.float32
                  )
                  for pred in preds
              ]
  
          tf_loss = tf.cast(
              lowering.export_to_tf_tensor(loss), tf.float32
          )
          if FLAGS.write_summary:
              tf_loss = _add_summary(
                  lowering, train_or_eval, tf_loss, scalars, global_step
              )
          master_to_slice_hook = mtf.MtfRestoreHook(lowering)
  
          if train_or_eval == "train":
              with mtf.utils.outside_all_rewrites():
                  saver = tf.train.Saver(
                      tf.global_variables(), save_relative_paths=True
                  )
                  tf.add_to_collection(tf.GraphKeys.SAVERS, saver)
                  saver_listener = mtf.MtfCheckpointSaverListener(
                      lowering
                  )
                  slice_to_master_hook = tf.train.CheckpointSaverHook(
                      FLAGS.checkpoint_dir,
                      save_steps=FLAGS.save_checkpoints_steps,
                      saver=saver,
                      listeners=[saver_listener],
                  )
                  captured_hooks.capture(
                      [master_to_slice_hook, slice_to_master_hook]
                  )
                  return tf.group([tf_loss] + tf_update_ops)
  
          else:  # train_or_eval == 'eval':
              if FLAGS.use_tpu:
                  tf_preds.extend([tf_loss, global_step])
                  tf_preds_dtypes = [
                      tf_pred.dtype for tf_pred in tf_preds
                  ]
                  tf_preds_shapes = [
                      tf_pred.shape for tf_pred in tf_preds
                  ]
                  captured_hooks.capture([master_to_slice_hook, None])
                  captured_output_dtypes_shapes.capture(
                      [tf_preds_dtypes, tf_preds_shapes]
                  )
                  return tpu_ops.outfeed_enqueue_tuple(tf_preds)
  
              else:
                  tf_preds.extend([tf_loss, global_step])
                  captured_hooks.capture([master_to_slice_hook, None])
                  return tf_preds
  
      return _model_fn, captured_hooks, captured_output_dtypes_shapes
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import re
  import mesh_tensorflow
  import numpy
  import six
  from six.moves import range
  from six.moves import zip
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  from mesh_tensorflow.experimental import input_reader
  from mesh_tensorflow.experimental import unet
  from tensorflow.contrib import summary
  from tensorflow.contrib import tpu
  from tensorflow.contrib.tpu.python.tpu import device_assignment
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.platform import flags
  from tensorflow.python.tpu.ops import tpu_ops
label_desc: Use image node_id 1 for calling a global function with example usage: _get_model_fn(train_or_eval, mesh_context) and returns: _model_fn,  captured_hooks,  captured_output_dtypes_shapes, tf_loss, tf, tpu_ops, tf_preds, tf