code: def model_fn(features, labels, mode, params):
      del params  # unused
      with variable_scope.variable_scope(
          "m", reuse=variable_scope.AUTO_REUSE
      ):
          w = variable_scope.get_variable("W", shape=[1000, 10])
      logits = math_ops.matmul(features, w)
      loss = losses.sparse_softmax_cross_entropy(
          labels=labels, logits=logits
      )
  
      if mode == model_fn_lib.ModeKeys.TRAIN:
          optimizer = training.RMSPropOptimizer(learning_rate=0.01)
          optimizer = tpu_optimizer.CrossShardOptimizer(optimizer)
          train_op = optimizer.minimize(
              loss, training.get_global_step()
          )
          return tpu_estimator.TPUEstimatorSpec(
              mode=model_fn_lib.ModeKeys.TRAIN,
              loss=loss,
              train_op=train_op,
          )
      elif mode == model_fn_lib.ModeKeys.EVAL:
  
          def metric_fn(labels, logits):
              labels = math_ops.cast(labels, dtypes.int64)
              logging.info("LABELS %s %s", labels, logits)
              return {
                  "recall@1": metrics_lib.recall_at_k(
                      labels, logits, 1
                  ),
                  "recall@5": metrics_lib.recall_at_k(
                      labels, logits, 5
                  ),
              }
  
          loss = losses.sparse_softmax_cross_entropy(
              labels=labels, logits=logits
          )
          eval_metrics = (metric_fn, [labels, logits])
          return tpu_estimator.TPUEstimatorSpec(
              mode=model_fn_lib.ModeKeys.EVAL,
              loss=loss,
              eval_metrics=eval_metrics,
          )
  
imports: import os
  import numpy
  from tensorflow.core.framework import summary_pb2
  from tensorflow.python.compat import v2_compat
  from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.lib.io import file_io
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import metrics
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.ops.losses import losses
  from tensorflow.python.platform import flags
  from tensorflow.python.platform import test
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.saved_model.pywrap_saved_model import metrics
  from tensorflow.python.tpu import async_checkpoint
  from tensorflow.python.tpu import tpu_optimizer
  from tensorflow.python.training import basic_session_run_hooks
  from tensorflow.python.training import training
  from tensorflow_estimator.python.estimator import estimator
  from tensorflow_estimator.python.estimator import model_fn
  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: model_fn(features, labels, mode, params) and returns: tpu_estimator, tpu_estimator, dict