code: def model_fn_builder(
      bert_config,
      init_checkpoint,
      learning_rate,
      num_train_steps,
      num_warmup_steps,
      use_tpu,
  ):
      """Returns `model_fn` closure for TPUEstimator."""
  
      def model_fn(
          features, labels, mode, params
      ):  # pylint: disable=unused-argument
          """The `model_fn` for TPUEstimator."""
  
          tf.logging.info("*** Features ***")
          for name in sorted(features.keys()):
              tf.logging.info(
                  "  name = %s, shape = %s"
                  % (name, features[name].shape)
              )
  
          # MTF setup.
          graph = mtf.Graph()
          mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
          layout_rules = mtf.convert_to_layout_rules(FLAGS.layout)
  
          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,
          )
          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,
          )
          mesh = mtf.Mesh(graph, "bert_mesh", var_placer)
  
          unique_ids = features["unique_ids"]
          input_ids = features["input_ids"]
          input_mask = features["input_mask"]
          segment_ids = features["segment_ids"]
  
          batch_size = input_ids.get_shape()[0].value
          batch_dim = mtf.Dimension("batch", batch_size)
          seq_length = input_ids.get_shape()[1].value
          seq_dim = mtf.Dimension("seq", seq_length)
  
          mtf_input_ids = mtf.import_tf_tensor(
              mesh, input_ids, [batch_dim, seq_dim]
          )
          mtf_input_mask = mtf.import_tf_tensor(
              mesh, input_mask, [batch_dim, seq_dim]
          )
          mtf_segment_ids = mtf.import_tf_tensor(
              mesh, segment_ids, [batch_dim, seq_dim]
          )
  
          is_training = mode == tf_estimator.ModeKeys.TRAIN
  
          (start_logits, end_logits) = create_model(
              bert_config=bert_config,
              is_training=is_training,
              input_ids=mtf_input_ids,
              input_mask=mtf_input_mask,
              segment_ids=mtf_segment_ids,
          )
  
          if mode == tf_estimator.ModeKeys.TRAIN:
  
              def compute_loss(logits, positions):
                  one_hot_positions = mtf.one_hot(
                      positions, output_dim=seq_dim
                  )
                  log_probs = mtf.log_softmax(logits, seq_dim)
                  loss = -mtf.reduce_mean(
                      mtf.reduce_sum(
                          one_hot_positions * log_probs,
                          reduced_dim=seq_dim,
                      )
                  )
                  return loss
  
              start_positions = features["start_positions"]
              mtf_start_positions = mtf.import_tf_tensor(
                  mesh, start_positions, [batch_dim]
              )
              end_positions = features["end_positions"]
              mtf_end_positions = mtf.import_tf_tensor(
                  mesh, end_positions, [batch_dim]
              )
  
              start_loss = compute_loss(
                  start_logits, mtf_start_positions
              )
              end_loss = compute_loss(end_logits, mtf_end_positions)
  
              total_loss = (start_loss + end_loss) / 2.0
              _, update_ops = optimization_lib.create_optimizer(
                  total_loss,
                  learning_rate,
                  num_train_steps,
                  num_warmup_steps,
                  max_optimized_variable_size=FLAGS.max_optimized_variable_size,
                  optimizer=FLAGS.optimizer,
                  clip_gradients=FLAGS.clip_gradients,
              )
          elif mode == tf_estimator.ModeKeys.PREDICT:
              start_logits = mtf.anonymize(start_logits)
              end_logits = mtf.anonymize(end_logits)
  
          lowering = mtf.Lowering(graph, {mesh: mesh_impl})
  
          if mode == tf_estimator.ModeKeys.TRAIN:
              tf_loss = tf.to_float(
                  lowering.export_to_tf_tensor(total_loss)
              )
              global_step = tf.train.get_global_step()
              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)
  
          tvars = tf.trainable_variables()
          initialized_variable_names = {}
          scaffold_fn = None
          if init_checkpoint:
              (
                  assignment_map,
                  initialized_variable_names,
              ) = bert_lib.get_assignment_map_from_checkpoint(
                  tvars, init_checkpoint
              )
              if use_tpu:
  
                  def tpu_scaffold():
                      tf.train.init_from_checkpoint(
                          init_checkpoint, assignment_map
                      )
                      return tf.train.Scaffold()
  
                  scaffold_fn = tpu_scaffold
              else:
                  tf.train.init_from_checkpoint(
                      init_checkpoint, assignment_map
                  )
  
          tf.logging.info("**** Trainable Variables ****")
          for var in tvars:
              init_string = ""
              if var.name in initialized_variable_names:
                  init_string = ", *INIT_FROM_CKPT*"
              tf.logging.info(
                  "  name = %s, shape = %s%s",
                  var.name,
                  var.shape,
                  init_string,
              )
  
          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.output_dir,
                      save_steps=1000,
                      saver=saver,
                      listeners=[saver_listener],
                  )
  
                  return tf_estimator.tpu.TPUEstimatorSpec(
                      mode,
                      loss=tf_loss,
                      train_op=train_op,
                      training_hooks=[restore_hook, saver_hook],
                      scaffold_fn=scaffold_fn,
                  )
              elif mode == tf_estimator.ModeKeys.PREDICT:
                  predictions = {
                      "unique_ids": unique_ids,
                      "start_logits": lowering.export_to_tf_tensor(
                          start_logits
                      ),
                      "end_logits": lowering.export_to_tf_tensor(
                          end_logits
                      ),
                  }
  
                  return tf_estimator.tpu.TPUEstimatorSpec(
                      mode,
                      prediction_hooks=[restore_hook],
                      predictions=predictions,
                      scaffold_fn=scaffold_fn,
                  )
              else:
                  raise ValueError(
                      "Only TRAIN and PREDICT modes are supported: %s"
                      % (mode)
                  )
  
      return model_fn
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import json
  import math
  import os
  import random
  import mesh_tensorflow
  import mesh_tensorflow.bert.bert
  import mesh_tensorflow.bert.optimization
  import mesh_tensorflow.bert.tokenization
  import six
  from six.moves import range
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 6 for calling a global function with example usage: model_fn_builder(bert_config, init_checkpoint, learning_rate, num_train_steps, num_warmup_steps, use_tpu) and returns: model_fn, loss, tf_estimator, tf, tf_estimator