code: def _train_phase(mesh_context, config, master):
      """Handles input pipeline and trains the network."""
      if FLAGS.num_train_iterations_per_loop <= 0:
          return
  
      def _run_train_phase():
          """The real function that runs the training phase."""
          # Setup input pipeline.
          ds_creator = unet.get_dataset_creator("train")
          mtf_shapes = unet.get_input_mtf_shapes("train")
  
          model_train_fn, train_hooks, _ = _get_model_fn(
              "train", mesh_context
          )
  
          if FLAGS.use_tpu:
              assert mesh_context.device_assignment
              assert mesh_context.num_cores
              simd_input_reader = input_reader.SimdMeshImplInputReader(
                  mesh_context.mesh_impl,
                  ds_creator,
                  mtf_shapes,
                  external_worker=(not FLAGS.on_gcp),
                  is_eval_mode=False,
              )
              train_computation = tpu.replicate(
                  computation=model_train_fn,
                  inputs=[[]] * mesh_context.num_cores,
                  infeed_queue=simd_input_reader.infeed_queue,
                  device_assignment=mesh_context.device_assignment,
              )
  
          else:
              placement_input_reader = (
                  input_reader.PlacementMeshImplInputReader(
                      mesh_context.mesh_impl,
                      ds_creator,
                      mtf_shapes,
                      is_eval_mode=False,
                  )
              )
              train_computation = placement_input_reader.gpu_placement(
                  model_train_fn
              )
  
          ###########################################################
          # Training.
          master_to_slice_hook, slice_to_master_hook = train_hooks.get()
          ckpt_loader_hook = _CkptLoaderHook()
          step_counter_hook = tf.train.StepCounterHook(every_n_steps=10)
          all_hooks = [
              ckpt_loader_hook,
              master_to_slice_hook,
              slice_to_master_hook,
              step_counter_hook,
          ]
  
          if FLAGS.write_summary:
              flush_summary = contrib_summary.flush()
  
          with tf.train.MonitoredTrainingSession(
              master=master,
              scaffold=_get_scaffold(additional_initializers=[]),
              hooks=all_hooks,
              config=config,
          ) as sess:
              if FLAGS.write_summary:
                  contrib_summary.initialize(session=sess)
  
              if FLAGS.use_tpu:
                  simd_input_reader.start_infeed_thread(
                      sess, FLAGS.num_train_iterations_per_loop
                  )
              else:
                  placement_input_reader.initialize(sess)
  
              for step in range(FLAGS.num_train_iterations_per_loop):
                  sess.run(train_computation)
                  if FLAGS.write_summary:
                      sess.run(flush_summary)
                  tf.logging.info("train steps: {}".format(step))
  
      with tf.Graph().as_default():
          if FLAGS.write_summary:
              summary_writer = contrib_summary.create_file_writer(
                  FLAGS.summary_dir
              )
              with summary_writer.as_default(), (
                  contrib_summary.always_record_summaries()
              ):
                  _run_train_phase()
          else:
              _run_train_phase()
  
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 4 for calling a global function with example usage: _train_phase(mesh_context, config, master) without return types