code: def create_estimator(
      model_name,
      hparams,
      run_config,
      schedule="train_and_evaluate",
      decode_hparams=None,
      use_tpu=False,
      use_tpu_estimator=False,
      use_xla=False,
      export_saved_model_api_version=1,
      use_guarantee_const_getter=False,
  ):
      """Create a T2T Estimator."""
      model_fn = t2t_model.T2TModel.make_estimator_model_fn(
          model_name,
          hparams,
          decode_hparams=decode_hparams,
          use_tpu=use_tpu,
      )
  
      del use_xla
      if use_tpu or use_tpu_estimator:
          from tensorflow.contrib.tpu.python.tpu import (
              tpu_estimator,
          )  # pylint: disable=g-import-not-at-top
  
          problem = hparams.problem
          batch_size = (
              problem.tpu_batch_size_per_shard(hparams)
              * run_config.tpu_config.num_shards
          )
          mlperf_log.transformer_print(
              key=mlperf_log.INPUT_BATCH_SIZE, value=batch_size
          )
          if getattr(hparams, "mtf_mode", False):
              batch_size = problem.tpu_batch_size_per_shard(hparams)
          predict_batch_size = batch_size
          if decode_hparams and decode_hparams.batch_size:
              predict_batch_size = decode_hparams.batch_size
          if decode_hparams and run_config.tpu_config:
              decode_hparams.add_hparam(
                  "iterations_per_loop",
                  run_config.tpu_config.iterations_per_loop,
              )
          if export_saved_model_api_version == 1:
              api_version_enum_name = (
                  tpu_estimator.ExportSavedModelApiVersion.V1
              )
              estimator_model_fn = model_fn
          elif export_saved_model_api_version == 2:
              api_version_enum_name = (
                  tpu_estimator.ExportSavedModelApiVersion.V2
              )
  
              def maybe_use_guarantee_const_getter_model_fn(
                  features, labels, mode, params
              ):
                  """Wrapper model_fn with guarantee_const getter."""
                  if not use_guarantee_const_getter:
                      return model_fn(features, labels, mode, params)
  
                  # It marks all weights as constant, which may improves TPU inference
                  # performance because it prevents the weights being transferred to the
                  # TPU. It will increase HBM "program" usage and reduce HBM "arguments"
                  # usage during TPU model serving.
                  def guarantee_const_getter(
                      getter, name, *args, **kwargs
                  ):
                      with tf.control_dependencies(None):
                          return tf.guarantee_const(
                              getter(name, *args, **kwargs),
                              name=name + "/GuaranteeConst",
                          )
  
                  @contextlib.contextmanager
                  def guarantee_const_scope():
                      var_scope = tf.get_variable_scope()
                      prev_custom_getter = var_scope.custom_getter
                      prev_caching_device = var_scope.caching_device
                      var_scope.set_custom_getter(
                          guarantee_const_getter
                      )
                      var_scope.set_caching_device(lambda op: op.device)
                      yield
                      var_scope.set_custom_getter(prev_custom_getter)
                      var_scope.set_caching_device(prev_caching_device)
  
                  with guarantee_const_scope():
                      return model_fn(features, labels, mode, params)
  
              def tpu_model_fn(features, labels, mode, params):
                  """Wrapper model_fn with tpu.rewrite / TPUPartitionedCall."""
                  if (
                      mode == tf_estimator.ModeKeys.PREDICT
                      and params["use_tpu"]
                  ):
                      batch_config = tpu_estimator.BatchConfig(
                          num_batch_threads=2,
                          max_batch_size=predict_batch_size,
                          batch_timeout_micros=60 * 1000,
                          allowed_batch_sizes=[predict_batch_size],
                      )
                      return tpu_estimator.model_fn_inference_on_tpu(
                          maybe_use_guarantee_const_getter_model_fn,
                          features=features,
                          labels=labels,
                          config=None,
                          params=params,
                          batch_config=batch_config,
                      )
                  else:
                      return model_fn(features, labels, mode, params)
  
              estimator_model_fn = tpu_model_fn
          else:
              raise ValueError(
                  "Flag export_saved_model_api_version must be 1 or 2."
              )
          estimator = contrib.tpu().TPUEstimator(
              model_fn=estimator_model_fn,
              model_dir=run_config.model_dir,
              config=run_config,
              use_tpu=use_tpu,
              train_batch_size=batch_size,
              eval_batch_size=batch_size
              if "eval" in schedule
              else None,
              predict_batch_size=predict_batch_size,
              export_saved_model_api_version=api_version_enum_name,
          )
      else:
          estimator = tf_estimator.Estimator(
              model_fn=model_fn,
              model_dir=run_config.model_dir,
              config=run_config,
          )
      return estimator
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import contextlib
  import json
  import os
  import random
  import numpy
  from tensor2tensor.utils import contrib
  from tensor2tensor.utils import decoding
  from tensor2tensor.utils import devices
  from tensor2tensor.utils import hparams_lib
  from tensor2tensor.utils import metrics_hook
  from tensor2tensor.utils import mlperf_log
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import t2t_model
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  from tensorflow.core.protobuf import rewriter_config_pb2
  from tensorflow.python import debug
label_desc: Use image node_id 6 for calling a global function with example usage: create_estimator(model_name, hparams, run_config, schedule, decode_hparams, use_tpu, use_tpu_estimator, use_xla, export_saved_model_api_version, use_guarantee_const_getter) and returns: estimator, model_fn, model_fn, tpu_estimator, model_fn, tf