code: def run_toy_model_tpu():
      """Run a toy model on TPU."""
      tpu_cluster_resolver = (
          tf.distribute.cluster_resolver.TPUClusterResolver(
              FLAGS.tpu, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project
          )
      )
  
      iterations_per_loop = FLAGS.iterations
      mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
      config = tpu_config.RunConfig(
          cluster=tpu_cluster_resolver,
          model_dir=FLAGS.model_dir,
          save_checkpoints_steps=None,  # Disable the default saver
          save_checkpoints_secs=None,  # Disable the default saver
          log_step_count_steps=iterations_per_loop,
          save_summary_steps=iterations_per_loop,
          tpu_config=tpu_config.TPUConfig(
              num_shards=mesh_shape.size,
              iterations_per_loop=iterations_per_loop,
              num_cores_per_replica=1,
              per_host_input_for_training=tpu_config.InputPipelineConfig.BROADCAST,
          ),
      )
      classifier = tpu_estimator.TPUEstimator(
          use_tpu=True,
          model_fn=model_fn,
          config=config,
          train_batch_size=FLAGS.batch_size,
          eval_batch_size=FLAGS.batch_size,
      )
      current_step = (
          estimator_lib._load_global_step_from_checkpoint_dir(
              FLAGS.model_dir
          )
      )  # pylint: disable=protected-access,line-too-long
      logging.info("Current step %d", current_step)
      if FLAGS.steps_per_checkpoint == 0:
          classifier.train(
              input_fn=ToyModelInput(), max_steps=FLAGS.train_steps
          )
          return
      while current_step < FLAGS.train_steps:
          next_checkpoint = min(
              current_step + FLAGS.steps_per_checkpoint,
              FLAGS.train_steps,
          )
          classifier.train(
              input_fn=ToyModelInput(), max_steps=next_checkpoint
          )
          current_step = next_checkpoint
          logging.info("Starting to evaluate.")
          eval_results = classifier.evaluate(
              input_fn=ToyModelInput(), steps=156
          )  # since we have 10000 examples and batch_size = 64 per host
          logging.info("Eval results: %s", eval_results)
  
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 3 for calling a global function with example usage: run_toy_model_tpu() without return types