docs: def model_fn(features, labels, mode, params):
      """
      Create the model for estimator api
  
      Args:
          features: tensor with shape
              [BATCH_SIZE, go.N, go.N, features_lib.NEW_FEATURES_PLANES]
          labels: dict from string to tensor with shape
              'pi_tensor': [BATCH_SIZE, go.N * go.N + 1]
              'value_tensor': [BATCH_SIZE]
          mode: a tf.estimator.ModeKeys (batchnorm params update for TRAIN only)
          params: A dictionary (Typically derived from the FLAGS object.)
      Returns: tf.estimator.EstimatorSpec with props
          mode: same as mode arg
          predictions: dict of tensors
              'policy': [BATCH_SIZE, go.N * go.N + 1]
              'value': [BATCH_SIZE]
          loss: a single value tensor
          train_op: train op
          eval_metric_ops
      return dict of tensors
          logits: [BATCH_SIZE, go.N * go.N + 1]
      """
  
      policy_output, value_output, logits = model_inference_fn(
          features, mode == tf.estimator.ModeKeys.TRAIN, params)
  
      # train ops
      policy_cost = tf.reduce_mean(
          tf.nn.softmax_cross_entropy_with_logits_v2(
              logits=logits, labels=tf.stop_gradient(labels['pi_tensor'])))
  
      value_cost = params['value_cost_weight'] * tf.reduce_mean(
          tf.square(value_output - labels['value_tensor']))
  
      reg_vars = [v for v in tf.trainable_variables()
                  if 'bias' not in v.name and 'beta' not in v.name]
      l2_cost = params['l2_strength'] * \
          tf.add_n([tf.nn.l2_loss(v) for v in reg_vars])
  
      combined_cost = policy_cost + value_cost + l2_cost
  
      global_step = tf.train.get_or_create_global_step()
      learning_rate = tf.train.piecewise_constant(
          global_step, params['lr_boundaries'], params['lr_rates'])
      update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
  
      # Insert quantization ops if requested
      if params['quantize']:
          if mode == tf.estimator.ModeKeys.TRAIN:
              tf.contrib.quantize.create_training_graph(
                  quant_delay=params['quant_delay'])
          else:
              tf.contrib.quantize.create_eval_graph()
  
      optimizer = tf.train.MomentumOptimizer(
          learning_rate, params['sgd_momentum'])
      if params['use_tpu']:
          optimizer = tpu_optimizer.CrossShardOptimizer(optimizer)
      with tf.control_dependencies(update_ops):
          train_op = optimizer.minimize(combined_cost, global_step=global_step)
  
      # Computations to be executed on CPU, outside of the main TPU queues.
      def eval_metrics_host_call_fn(policy_output, value_output, pi_tensor, policy_cost,
                                    value_cost, l2_cost, combined_cost, step,
                                    est_mode=tf.estimator.ModeKeys.TRAIN):
          policy_entropy = -tf.reduce_mean(tf.reduce_sum(
              policy_output * tf.log(policy_output), axis=1))
          # pi_tensor is one_hot when generated from sgfs (for supervised learning)
          # and soft-max when using self-play records. argmax normalizes the two.
          policy_target_top_1 = tf.argmax(pi_tensor, axis=1)
  
          policy_output_in_top1 = tf.to_float(
              tf.nn.in_top_k(policy_output, policy_target_top_1, k=1))
          policy_output_in_top3 = tf.to_float(
              tf.nn.in_top_k(policy_output, policy_target_top_1, k=3))
  
          policy_top_1_confidence = tf.reduce_max(policy_output, axis=1)
          policy_target_top_1_confidence = tf.boolean_mask(
              policy_output,
              tf.one_hot(policy_target_top_1, tf.shape(policy_output)[1]))
  
          value_cost_normalized = value_cost / params['value_cost_weight']
  
          with tf.variable_scope("metrics"):
              metric_ops = {
                  'policy_cost': tf.metrics.mean(policy_cost),
                  'value_cost': tf.metrics.mean(value_cost),
                  'value_cost_normalized': tf.metrics.mean(value_cost_normalized),
                  'l2_cost': tf.metrics.mean(l2_cost),
                  'policy_entropy': tf.metrics.mean(policy_entropy),
                  'combined_cost': tf.metrics.mean(combined_cost),
  
                  'policy_accuracy_top_1': tf.metrics.mean(policy_output_in_top1),
                  'policy_accuracy_top_3': tf.metrics.mean(policy_output_in_top3),
                  'policy_top_1_confidence': tf.metrics.mean(policy_top_1_confidence),
                  'policy_target_top_1_confidence': tf.metrics.mean(
                      policy_target_top_1_confidence),
                  'value_confidence': tf.metrics.mean(tf.abs(value_output)),
              }
  
          if est_mode == tf.estimator.ModeKeys.EVAL:
              return metric_ops
  
          # NOTE: global_step is rounded to a multiple of FLAGS.summary_steps.
          eval_step = tf.reduce_min(step)
  
          # Create summary ops so that they show up in SUMMARIES collection
          # That way, they get logged automatically during training
          summary_writer = summary.create_file_writer(FLAGS.work_dir)
          with summary_writer.as_default(), \
                  summary.record_summaries_every_n_global_steps(
                      params['summary_steps'], eval_step):
              for metric_name, metric_op in metric_ops.items():
                  summary.scalar(metric_name, metric_op[1], step=eval_step)
  
          # Reset metrics occasionally so that they are mean of recent batches.
          reset_op = tf.variables_initializer(tf.local_variables("metrics"))
          cond_reset_op = tf.cond(
              tf.equal(eval_step % params['summary_steps'], tf.to_int64(1)),
              lambda: reset_op,
              lambda: tf.no_op())
  
          return summary.all_summary_ops() + [cond_reset_op]
  
      metric_args = [
          policy_output,
          value_output,
          labels['pi_tensor'],
          tf.reshape(policy_cost, [1]),
          tf.reshape(value_cost, [1]),
          tf.reshape(l2_cost, [1]),
          tf.reshape(combined_cost, [1]),
          tf.reshape(global_step, [1]),
      ]
  
      predictions = {
          'policy_output': policy_output,
          'value_output': value_output,
      }
  
      eval_metrics_only_fn = functools.partial(
          eval_metrics_host_call_fn, est_mode=tf.estimator.ModeKeys.EVAL)
      host_call_fn = functools.partial(
          eval_metrics_host_call_fn, est_mode=tf.estimator.ModeKeys.TRAIN)
  
      tpu_estimator_spec = tpu_estimator.TPUEstimatorSpec(
          mode=mode,
          predictions=predictions,
          loss=combined_cost,
          train_op=train_op,
          eval_metrics=(eval_metrics_only_fn, metric_args),
          host_call=(host_call_fn, metric_args)
      )
      if params['use_tpu']:
          return tpu_estimator_spec
      else:
          return tpu_estimator_spec.as_estimator_spec()