def _get_eval_estimator_spec(gan_model,
                             gan_loss,
                             get_eval_metric_ops_fn=None,
                             name=None):
  """Return an EstimatorSpec for the eval case."""
  scalar_loss = gan_loss.generator_loss + gan_loss.discriminator_loss
  with ops.name_scope(None, 'metrics',
                      [gan_loss.generator_loss, gan_loss.discriminator_loss]):

    def _summary_key(head_name, val):
      return '%s/%s' % (val, head_name) if head_name else val

    eval_metric_ops = {
        _summary_key(name, 'generator_loss'):
            metrics_lib.mean(gan_loss.generator_loss),
        _summary_key(name, 'discriminator_loss'):
            metrics_lib.mean(gan_loss.discriminator_loss)
    }
    if get_eval_metric_ops_fn is not None:
      custom_eval_metric_ops = get_eval_metric_ops_fn(gan_model)
      if not isinstance(custom_eval_metric_ops, dict):
        raise TypeError('get_eval_metric_ops_fn must return a dict, '
                        'received: {}'.format(custom_eval_metric_ops))
      eval_metric_ops.update(custom_eval_metric_ops)
  return model_fn_lib.EstimatorSpec(
      mode=model_fn_lib.ModeKeys.EVAL,
      predictions=gan_model.generated_data,
      loss=scalar_loss,
      eval_metric_ops=eval_metric_ops)

def _convert_keras_metrics_to_estimator(model):
  """Convert metrics from a Keras model to ops used by the Estimator framework.

  Args:
    model: A `tf.keras.Model` object.

  Returns:
    Dictionary mapping metric names to tuples of (value, update) ops. May return
    `None` if the model does not contain any metrics.
  """
  if not getattr(model, 'metrics', None):
    return None

  # TODO(psv/fchollet): support stateful metrics
  eval_metric_ops = {}
  # When each metric maps to an output
  if isinstance(model.metrics, dict):
    for i, output_name in enumerate(model.metrics.keys()):
      metric_name = model.metrics[output_name]
      if callable(metric_name):
        metric_name = metric_name.__name__
      # When some outputs use the same metric
      if list(model.metrics.values()).count(metric_name) > 1:
        metric_name += '_' + output_name
      eval_metric_ops[metric_name] = metrics_module.mean(
          model.metrics_tensors[i - len(model.metrics)])
  else:
    for i, metric_name in enumerate(model.metrics):
      if callable(metric_name):
        metric_name = metric_name.__name__
      eval_metric_ops[metric_name] = metrics_module.mean(
          model.metrics_tensors[i])
  return eval_metric_ops

 def _eval_metric_ops(
     self, labels, probabilities, weights, unreduced_loss,
     regularization_loss):
   """Returns a dict of metrics for eval_metric_ops."""
   with ops.name_scope(
       None, 'metrics',
       [labels, probabilities, weights, unreduced_loss, regularization_loss]):
     keys = metric_keys.MetricKeys
     metric_ops = {
         # Estimator already adds a metric for loss.
         head_lib._summary_key(self._name, keys.LOSS_MEAN):  # pylint:disable=protected-access
             metrics_lib.mean(
                 values=unreduced_loss,
                 weights=weights,
                 name=keys.LOSS_MEAN),
         head_lib._summary_key(self._name, keys.AUC):  # pylint:disable=protected-access
             metrics_lib.auc(labels=labels, predictions=probabilities,
                             weights=weights, name=keys.AUC),
         head_lib._summary_key(self._name, keys.AUC_PR):  # pylint:disable=protected-access
             metrics_lib.auc(labels=labels, predictions=probabilities,
                             weights=weights, curve='PR',
                             name=keys.AUC_PR),
     }
     if regularization_loss is not None:
       loss_regularization_key = head_lib._summary_key(  # pylint:disable=protected-access
           self._name, keys.LOSS_REGULARIZATION)
       metric_ops[loss_regularization_key] = (
           metrics_lib.mean(
               values=regularization_loss,
               name=keys.LOSS_REGULARIZATION))
     for threshold in self._thresholds:
       accuracy_key = keys.ACCURACY_AT_THRESHOLD % threshold
       metric_ops[head_lib._summary_key(self._name, accuracy_key)] = (  # pylint:disable=protected-access
           head_lib._accuracy_at_threshold(  # pylint:disable=protected-access
               labels=labels,
               predictions=probabilities,
               weights=weights,
               threshold=threshold,
               name=accuracy_key))
       # Precision for positive examples.
       precision_key = keys.PRECISION_AT_THRESHOLD % threshold
       metric_ops[head_lib._summary_key(self._name, precision_key)] = (  # pylint:disable=protected-access
           head_lib._precision_at_threshold(  # pylint:disable=protected-access
               labels=labels,
               predictions=probabilities,
               weights=weights,
               threshold=threshold,
               name=precision_key))
       # Recall for positive examples.
       recall_key = keys.RECALL_AT_THRESHOLD % threshold
       metric_ops[head_lib._summary_key(self._name, recall_key)] = (  # pylint:disable=protected-access
           head_lib._recall_at_threshold(  # pylint:disable=protected-access
               labels=labels,
               predictions=probabilities,
               weights=weights,
               threshold=threshold,
               name=recall_key))
   return metric_ops

  def model_fn(features, labels, mode):
    """model_fn for keras Estimator."""
    model = _clone_and_build_model(mode, keras_model, custom_objects, features,
                                   labels)
    # Get inputs to EstimatorSpec
    predictions = dict(zip(model.output_names, model.outputs))

    loss = None
    train_op = None
    eval_metric_ops = None

    # Set loss and metric only during train and evaluate.
    if mode is not model_fn_lib.ModeKeys.PREDICT:
      if mode is model_fn_lib.ModeKeys.TRAIN:
        model._make_train_function()  # pylint: disable=protected-access
      else:
        model._make_test_function()  # pylint: disable=protected-access
      loss = model.total_loss

      if model.metrics:
        # TODO(fchollet): support stateful metrics
        eval_metric_ops = {}
        # When each metric maps to an output
        if isinstance(model.metrics, dict):
          for i, output_name in enumerate(model.metrics.keys()):
            metric_name = model.metrics[output_name]
            if callable(metric_name):
              metric_name = metric_name.__name__
            # When some outputs use the same metric
            if list(model.metrics.values()).count(metric_name) > 1:
              metric_name += '_' + output_name
            eval_metric_ops[metric_name] = metrics_module.mean(
                model.metrics_tensors[i - len(model.metrics)])
        else:
          for i, metric_name in enumerate(model.metrics):
            if callable(metric_name):
              metric_name = metric_name.__name__
            eval_metric_ops[metric_name] = metrics_module.mean(
                model.metrics_tensors[i])

    # Set train_op only during train.
    if mode is model_fn_lib.ModeKeys.TRAIN:
      train_op = model.train_function.updates_op

    if not model._is_graph_network:
      # Reset model state to original state,
      # to avoid `model_fn` being destructive for the initial model argument.
      _in_place_subclassed_model_state_restoration(keras_model)
    return model_fn_lib.EstimatorSpec(
        mode=mode,
        predictions=predictions,
        loss=loss,
        train_op=train_op,
        eval_metric_ops=eval_metric_ops,
        export_outputs={
            _DEFAULT_SERVING_KEY:
            export_lib.export_output.PredictOutput(predictions)
        })

def _convert_keras_metrics_to_estimator(model):
  """Convert metrics from a Keras model to ops used by the Estimator framework.

  Args:
    model: A `tf.keras.Model` object.

  Returns:
    Dictionary mapping metric names to tuples of (value, update) ops. May return
    `None` if the model does not contain any metrics.
  """
  if not getattr(model, 'metrics', None):
    return None

  eval_metric_ops = {}

  def get_metric_name(metric):
    if isinstance(metric, metrics.Metric):
      return metric.name
    if callable(metric):
      return metric.__name__
    assert isinstance(metric, six.string_types)
    return metric

  # When each metric maps to an output
  if isinstance(model.metrics, dict):
    for i, output_name in enumerate(model.metrics.keys()):
      # `metric` is the user given metric value in `compile`. This can be
      # metric name (`acc`), metric function (binary_accuracy) or a metric
      # object (BinaryAccuracy()).
      metric = model.metrics[output_name]
      metric_name = get_metric_name(metric)
      # When some outputs use the same metric
      if list(model.metrics.values()).count(metric_name) > 1:
        metric_name += '_' + output_name
      if isinstance(metric, metrics.Metric):
        eval_metric_ops[metric_name] = metric
      else:
        eval_metric_ops[metric_name] = metrics_module.mean(
            model.metrics_tensors[i - len(model.metrics)])
  else:
    for i, metric in enumerate(model.metrics):
      metric_name = get_metric_name(metric)
      if isinstance(metric, metrics.Metric):
        eval_metric_ops[metric_name] = metric
      else:
        eval_metric_ops[metric_name] = metrics_module.mean(
            model.metrics_tensors[i])
  return eval_metric_ops

  def _merge_eval(self, all_estimator_spec):
    """Merges list of `EstimatorSpec` for eval.

    Args:
      all_estimator_spec: list of `EstimatorSpec` for the individual heads.

    Returns:
      `EstimatorSpec` that merges all heads for EVAL.
    """
    predictions = {}
    metrics = {}
    losses = []
    with ops.name_scope('merge_eval'):
      for head, spec in zip(self._heads, all_estimator_spec):
        losses.append(spec.loss)
        head_name = head.name
        # Loss metric is not added by default.
        loss_name = head_lib._summary_key(  # pylint:disable=protected-access
            head_name, metric_keys.MetricKeys.LOSS)
        metrics[loss_name] = metrics_lib.mean(spec.loss, name=loss_name)
        # Metric keys already contain head.name.
        metrics.update(spec.eval_metric_ops or {})
        for k, v in six.iteritems(spec.predictions):
          predictions[(head_name, k)] = v
      loss = _merge_losses(losses, self._head_weights)

    return model_fn.EstimatorSpec(
        mode=model_fn.ModeKeys.EVAL,
        predictions=predictions,
        loss=loss,
        eval_metric_ops=metrics)

def _predictions_mean(predictions, weights=None, name=None):
  with ops.name_scope(
      name, 'predictions_mean', (predictions, weights)) as scope:
    predictions = math_ops.to_float(predictions, name='predictions')
    if weights is not None:
      weights = weights_broadcast_ops.broadcast_weights(weights, predictions)
    return metrics_lib.mean(predictions, weights=weights, name=scope)

 def _eval_metric_ops(self, labels, class_ids, weights, weighted_sum_loss,
                      example_weight_sum):
   """Returns the Eval metric ops."""
   with ops.name_scope(
       None, 'metrics',
       (labels, class_ids, weights, weighted_sum_loss, example_weight_sum)):
     keys = metric_keys.MetricKeys
     metric_ops = {
         # Estimator already adds a metric for loss.
         # TODO(xiejw): Any other metrics?
         _summary_key(self._name, keys.LOSS_MEAN):
             metrics_lib.mean(
                 # Both values and weights here are reduced, scalar Tensors.
                 # values is the actual mean we want -- weights represents the
                 # total weight of the batch and is needed to calculate
                 # update_op over many batches.
                 values=(weighted_sum_loss / example_weight_sum),
                 weights=example_weight_sum,
                 name=keys.LOSS_MEAN),
         _summary_key(self._name, keys.ACCURACY):
             metrics_lib.accuracy(
                 labels=labels,
                 predictions=class_ids,
                 weights=weights,
                 name=keys.ACCURACY),
     }
   return metric_ops

 def metric_fn(
     generator_inputs, generated_data, real_data, discriminator_real_outputs,
     discriminator_gen_outputs, generator_loss, discriminator_loss):
   """`metric_fn` used in TPUEstimator to calculate metrics."""
   eval_metric_ops = {
       'generator_loss': metrics_lib.mean(generator_loss),
       'discriminator_loss': metrics_lib.mean(discriminator_loss),
   }
   custom_eval_metric_ops = get_eval_metric_ops_fn(
       generator_inputs, generated_data, real_data,
       discriminator_real_outputs, discriminator_gen_outputs)
   if not isinstance(custom_eval_metric_ops, dict):
     raise TypeError('`get_eval_metric_ops_fn` must return a dict, '
                     'received: {}'.format(custom_eval_metric_ops))
   eval_metric_ops.update(custom_eval_metric_ops)
   return eval_metric_ops

def _sigmoid_entropy(probabilities, targets, weights=None):
  return metrics.mean(
      losses.sigmoid_cross_entropy(probabilities,
                                   _squeeze_and_onehot(
                                       targets,
                                       array_ops.shape(probabilities)[1])),
      weights=weights)

 def model_fn(features, labels, mode):
   _ = labels
   step = training.get_global_step()
   w = variable_scope.get_variable(
       'w',
       shape=[],
       initializer=init_ops.zeros_initializer(),
       dtype=dtypes.int64)
   if estimator_lib.ModeKeys.TRAIN == mode:
     # to consume features, we have control dependency
     with ops.control_dependencies([features]):
       step_inc = state_ops.assign_add(training.get_global_step(), 1)
     with ops.control_dependencies([step_inc]):
       assign_w_to_step_plus_2 = w.assign(step + 2)
     return estimator_lib.EstimatorSpec(
         mode,
         loss=constant_op.constant(3.),
         train_op=assign_w_to_step_plus_2)
   if estimator_lib.ModeKeys.EVAL == mode:
     # to consume features, we have control dependency
     with ops.control_dependencies([features]):
       loss = constant_op.constant(5.)
     return estimator_lib.EstimatorSpec(
         mode,
         loss=loss,
         # w is constant in each step, so the mean.
         # w = 0 if step==0 else step+2
         eval_metric_ops={'mean_of_const': metrics_lib.mean(w)})

  def _evaluate_model(self,
                      input_fn,
                      hooks=None,
                      checkpoint_path=None,
                      name=''):
    """Evaluates the model using the training.evaluation library."""
    # Check that model has been trained (if nothing has been set explicitly).
    if not checkpoint_path:
      latest_path = saver.latest_checkpoint(self._model_dir)
      if not latest_path:
        raise ValueError('Could not find trained model in model_dir: {}.'.
                         format(self._model_dir))
      checkpoint_path = latest_path

    # Setup output directory.
    eval_dir = os.path.join(self._model_dir, 'eval' if not name else
                            'eval_' + name)

    with ops.Graph().as_default() as g:
      random_seed.set_random_seed(self._config.tf_random_seed)
      global_step_tensor = self._create_and_assert_global_step(g)
      features, labels = self._get_features_and_labels_from_input_fn(
          input_fn, model_fn_lib.ModeKeys.EVAL)
      estimator_spec = self._call_model_fn(
          features, labels, model_fn_lib.ModeKeys.EVAL)

      if model_fn_lib.LOSS_METRIC_KEY in estimator_spec.eval_metric_ops:
        raise ValueError(
            'Metric with name "%s" is not allowed, because Estimator ' % (
                model_fn_lib.LOSS_METRIC_KEY) +
            'already defines a default metric with the same name.')
      estimator_spec.eval_metric_ops[
          model_fn_lib.LOSS_METRIC_KEY] = metrics_lib.mean(estimator_spec.loss)

      update_op, eval_dict = _extract_metric_update_ops(
          estimator_spec.eval_metric_ops)

      if ops.GraphKeys.GLOBAL_STEP in eval_dict:
        raise ValueError(
            'Metric with name `global_step` is not allowed, because Estimator '
            'already defines a default metric with the same name.')
      eval_dict[ops.GraphKeys.GLOBAL_STEP] = global_step_tensor

      eval_results = evaluation._evaluate_once(  # pylint: disable=protected-access
          checkpoint_path=checkpoint_path,
          master=self._config.evaluation_master,
          scaffold=estimator_spec.scaffold,
          eval_ops=update_op,
          final_ops=eval_dict,
          hooks=hooks,
          config=self._session_config)

      _write_dict_to_summary(
          output_dir=eval_dir,
          dictionary=eval_results,
          current_global_step=eval_results[ops.GraphKeys.GLOBAL_STEP])

    return eval_results

def _r2(probabilities, targets, weights=None):
  targets = math_ops.cast(targets, dtypes.float32)
  y_mean = math_ops.reduce_mean(targets, 0)
  squares_total = math_ops.reduce_sum(
      math_ops.squared_difference(targets, y_mean), 0)
  squares_residuals = math_ops.reduce_sum(
      math_ops.squared_difference(targets, probabilities), 0)
  score = 1 - math_ops.reduce_sum(squares_residuals / squares_total)
  return metrics.mean(score, weights=weights)

 def _eval_metric_ops(self, labels, probabilities, weights, weighted_sum_loss,
                      example_weight_sum):
   """Returns a dict of metrics for eval_metric_ops."""
   with ops.name_scope(
       None, 'metrics',
       [labels, probabilities, weights, weighted_sum_loss, example_weight_sum
       ]):
     keys = metric_keys.MetricKeys
     metric_ops = {
         # Estimator already adds a metric for loss.
         head_lib._summary_key(self._name, keys.LOSS_MEAN):  # pylint:disable=protected-access
             metrics_lib.mean(
                 # Both values and weights here are reduced, scalar Tensors.
                 # values is the actual mean we want, but we pass the scalar
                 # example_weight_sum in order to return the correct update_op
                 # alongside the value_op for streaming metrics.
                 values=(weighted_sum_loss / example_weight_sum),
                 weights=example_weight_sum,
                 name=keys.LOSS_MEAN),
         head_lib._summary_key(self._name, keys.AUC):  # pylint:disable=protected-access
             metrics_lib.auc(labels=labels, predictions=probabilities,
                             weights=weights, name=keys.AUC),
         head_lib._summary_key(self._name, keys.AUC_PR):  # pylint:disable=protected-access
             metrics_lib.auc(labels=labels, predictions=probabilities,
                             weights=weights, curve='PR',
                             name=keys.AUC_PR),
     }
     for threshold in self._thresholds:
       accuracy_key = keys.ACCURACY_AT_THRESHOLD % threshold
       metric_ops[head_lib._summary_key(self._name, accuracy_key)] = (  # pylint:disable=protected-access
           head_lib._accuracy_at_threshold(  # pylint:disable=protected-access
               labels=labels,
               predictions=probabilities,
               weights=weights,
               threshold=threshold,
               name=accuracy_key))
       # Precision for positive examples.
       precision_key = keys.PRECISION_AT_THRESHOLD % threshold
       metric_ops[head_lib._summary_key(self._name, precision_key)] = (  # pylint:disable=protected-access
           head_lib._precision_at_threshold(  # pylint:disable=protected-access
               labels=labels,
               predictions=probabilities,
               weights=weights,
               threshold=threshold,
               name=precision_key))
       # Recall for positive examples.
       recall_key = keys.RECALL_AT_THRESHOLD % threshold
       metric_ops[head_lib._summary_key(self._name, recall_key)] = (  # pylint:disable=protected-access
           head_lib._recall_at_threshold(  # pylint:disable=protected-access
               labels=labels,
               predictions=probabilities,
               weights=weights,
               threshold=threshold,
               name=recall_key))
   return metric_ops

  def create_estimator_spec(
      self, features, mode, logits, labels=None, train_op_fn=None):
    """See `Head`."""
    with variable_scope.variable_scope(
        None,
        default_name='regression_head',
        values=(tuple(six.itervalues(features)) + (labels, logits))):

      # Predict.
      logits = _check_logits(logits, self._logits_dimension)
      predictions = {prediction_keys.PredictionKeys.PREDICTIONS: logits}
      if mode == model_fn.ModeKeys.PREDICT:
        return model_fn.EstimatorSpec(
            mode=model_fn.ModeKeys.PREDICT,
            predictions=predictions,
            export_outputs={'': export_output.RegressionOutput(value=logits)})

      # Eval.
      labels = _check_labels(_maybe_expand_dim(math_ops.to_float(labels)),
                             self._logits_dimension)
      unweighted_loss = losses.mean_squared_error(
          labels=labels, predictions=logits, reduction=losses.Reduction.NONE)
      weights = (
          1. if (self._weight_feature_key is None) else
          features[self._weight_feature_key])
      weights = _maybe_expand_dim(math_ops.to_float(weights, name='weights'))
      training_loss = losses.compute_weighted_loss(
          unweighted_loss, weights=weights, reduction=losses.Reduction.SUM)
      if mode == model_fn.ModeKeys.EVAL:
        # Estimator already adds a metric for loss.
        eval_metric_ops = {
            metric_keys.MetricKeys.LOSS_MEAN: metrics_lib.mean(
                unweighted_loss, weights=weights)
        }
        return model_fn.EstimatorSpec(
            mode=model_fn.ModeKeys.EVAL,
            predictions=predictions,
            loss=training_loss,
            eval_metric_ops=eval_metric_ops)

      # Train.
      if train_op_fn is None:
        raise ValueError('train_op_fn can not be None.')
      logging_ops.scalar_summary(metric_keys.MetricKeys.LOSS, training_loss)
      logging_ops.scalar_summary(
          metric_keys.MetricKeys.LOSS_MEAN,
          losses.compute_weighted_loss(
              unweighted_loss, weights=weights,
              reduction=losses.Reduction.MEAN))
      return model_fn.EstimatorSpec(
          mode=model_fn.ModeKeys.TRAIN,
          predictions=predictions,
          loss=training_loss,
          train_op=train_op_fn(training_loss))

  def testMetricsCollection(self):

    def _enqueue_vector(sess, queue, values, shape=None):
      if not shape:
        shape = (1, len(values))
      dtype = queue.dtypes[0]
      sess.run(
          queue.enqueue(constant_op.constant(
              values, dtype=dtype, shape=shape)))

    meta_graph_filename = os.path.join(
        _TestDir("metrics_export"), "meta_graph.pb")

    graph = ops.Graph()
    with self.session(graph=graph) as sess:
      values_queue = data_flow_ops.FIFOQueue(
          4, dtypes.float32, shapes=(1, 2))
      _enqueue_vector(sess, values_queue, [0, 1])
      _enqueue_vector(sess, values_queue, [-4.2, 9.1])
      _enqueue_vector(sess, values_queue, [6.5, 0])
      _enqueue_vector(sess, values_queue, [-3.2, 4.0])
      values = values_queue.dequeue()

      _, update_op = metrics.mean(values)

      initializer = variables.local_variables_initializer()
      self.evaluate(initializer)
      self.evaluate(update_op)

    meta_graph.export_scoped_meta_graph(
        filename=meta_graph_filename, graph=graph)

    # Verifies that importing a meta_graph with LOCAL_VARIABLES collection
    # works correctly.
    graph = ops.Graph()
    with self.session(graph=graph) as sess:
      meta_graph.import_scoped_meta_graph(meta_graph_filename)
      initializer = variables.local_variables_initializer()
      self.evaluate(initializer)

    # Verifies that importing an old meta_graph where "local_variables"
    # collection is of node_list type works, but cannot build initializer
    # with the collection.
    graph = ops.Graph()
    with self.session(graph=graph) as sess:
      meta_graph.import_scoped_meta_graph(
          test.test_src_dir_path(
              "python/framework/testdata/metrics_export_meta_graph.pb"))
      self.assertEqual(len(ops.get_collection(ops.GraphKeys.LOCAL_VARIABLES)),
                       2)
      with self.assertRaisesRegexp(
          AttributeError, "'Tensor' object has no attribute 'initializer'"):
        initializer = variables.local_variables_initializer()

  def create_estimator_spec(
      self, features, mode, logits, labels=None, train_op_fn=None):
    """See `Head`."""
    # Predict.
    with ops.name_scope('head'):
      logits = _check_logits(logits, self._logits_dimension)
      predictions = {prediction_keys.PredictionKeys.PREDICTIONS: logits}
      if mode == model_fn.ModeKeys.PREDICT:
        regression_output = export_output.RegressionOutput(value=logits)
        return model_fn.EstimatorSpec(
            mode=model_fn.ModeKeys.PREDICT,
            predictions=predictions,
            export_outputs={
                _DEFAULT_SERVING_KEY: regression_output,
                _REGRESS_SERVING_KEY: regression_output,
                _PREDICT_SERVING_KEY: export_output.PredictOutput(predictions)
            })

      # Eval.
      unweighted_loss, _ = self.create_loss(
          features=features, mode=mode, logits=logits, labels=labels)
      weights = _weights(features, self._weight_column)
      training_loss = losses.compute_weighted_loss(
          unweighted_loss, weights=weights, reduction=losses.Reduction.SUM)
      if mode == model_fn.ModeKeys.EVAL:
        # Estimator already adds a metric for loss.
        eval_metric_ops = {
            metric_keys.MetricKeys.LOSS_MEAN: metrics_lib.mean(
                unweighted_loss, weights=weights)
        }
        return model_fn.EstimatorSpec(
            mode=model_fn.ModeKeys.EVAL,
            predictions=predictions,
            loss=training_loss,
            eval_metric_ops=eval_metric_ops)

      # Train.
      if train_op_fn is None:
        raise ValueError('train_op_fn can not be None.')
    with ops.name_scope(''):
      summary.scalar(
          _summary_key(self._name, metric_keys.MetricKeys.LOSS),
          training_loss)
      summary.scalar(
          _summary_key(self._name, metric_keys.MetricKeys.LOSS_MEAN),
          losses.compute_weighted_loss(
              unweighted_loss, weights=weights,
              reduction=losses.Reduction.MEAN))
    return model_fn.EstimatorSpec(
        mode=model_fn.ModeKeys.TRAIN,
        predictions=predictions,
        loss=training_loss,
        train_op=train_op_fn(training_loss))

  def _evaluate_model(self,
                      input_fn,
                      hooks=None,
                      checkpoint_path=None,
                      name=''):
    """Evaluates the model using the training.evaluation library."""
    # Check that model has been trained (if nothing has been set explicitly).
    if not checkpoint_path:
      latest_path = saver.latest_checkpoint(self._model_dir)
      if not latest_path:
        raise ValueError('Could not find trained model in model_dir: {}.'.
                         format(self._model_dir))
      checkpoint_path = latest_path

    # Setup output directory.
    eval_dir = os.path.join(self._model_dir, 'eval' if not name else
                            'eval_' + name)

    with ops.Graph().as_default() as g:
      random_seed.set_random_seed(self._config.tf_random_seed)
      global_step_tensor = training.create_global_step(g)
      features, labels = input_fn()
      estimator_spec = self._call_model_fn(
          features, labels, model_fn_lib.ModeKeys.EVAL)

      self._verify_default_metric_key(model_fn_lib.MetricKeys.LOSS,
                                      estimator_spec.eval_metric_ops)
      estimator_spec.eval_metric_ops[
          model_fn_lib.MetricKeys.LOSS] = metrics_lib.mean(estimator_spec.loss)

      update_op, eval_dict = _extract_metric_update_ops(
          estimator_spec.eval_metric_ops)

      self._verify_default_metric_key(ops.GraphKeys.GLOBAL_STEP, eval_dict)
      eval_dict[ops.GraphKeys.GLOBAL_STEP] = global_step_tensor

      eval_results = evaluation._evaluate_once(  # pylint: disable=protected-access
          checkpoint_path=checkpoint_path,
          master=self._config.evaluation_master,
          scaffold=estimator_spec.scaffold,
          eval_ops=update_op,
          final_ops=eval_dict,
          hooks=hooks,
          config=config_pb2.ConfigProto(allow_soft_placement=True))

      _write_dict_to_summary(
          output_dir=eval_dir,
          dictionary=eval_results,
          current_global_step=eval_results[ops.GraphKeys.GLOBAL_STEP])

    return eval_results

  def create_estimator_spec(
      self, features, mode, logits, labels=None, train_op_fn=None):
    """See `Head`."""
    # Predict.
    with ops.name_scope('head'):
      logits = _check_logits(logits, self._logits_dimension)
      predictions = {prediction_keys.PredictionKeys.PREDICTIONS: logits}
      if mode == model_fn.ModeKeys.PREDICT:
        return model_fn.EstimatorSpec(
            mode=model_fn.ModeKeys.PREDICT,
            predictions=predictions,
            export_outputs={'': export_output.RegressionOutput(value=logits)})

      # Eval.
      labels = _check_labels(_maybe_expand_dim(math_ops.to_float(labels)),
                             self._logits_dimension)
      unweighted_loss = losses.mean_squared_error(
          labels=labels, predictions=logits, reduction=losses.Reduction.NONE)
      weights = _weights(features, self._weight_column)
      training_loss = losses.compute_weighted_loss(
          unweighted_loss, weights=weights, reduction=losses.Reduction.SUM)
      if mode == model_fn.ModeKeys.EVAL:
        # Estimator already adds a metric for loss.
        eval_metric_ops = {
            metric_keys.MetricKeys.LOSS_MEAN: metrics_lib.mean(
                unweighted_loss, weights=weights)
        }
        return model_fn.EstimatorSpec(
            mode=model_fn.ModeKeys.EVAL,
            predictions=predictions,
            loss=training_loss,
            eval_metric_ops=eval_metric_ops)

      # Train.
      if train_op_fn is None:
        raise ValueError('train_op_fn can not be None.')
    with ops.name_scope(''):
      summary.scalar(metric_keys.MetricKeys.LOSS, training_loss)
      summary.scalar(metric_keys.MetricKeys.LOSS_MEAN,
                     losses.compute_weighted_loss(
                         unweighted_loss,
                         weights=weights,
                         reduction=losses.Reduction.MEAN))
    return model_fn.EstimatorSpec(
        mode=model_fn.ModeKeys.TRAIN,
        predictions=predictions,
        loss=training_loss,
        train_op=train_op_fn(training_loss))

 def test_metrics_consistent(self):
   # Tests that the identity metrics used to report in-sample predictions match
   # the behavior of standard metrics.
   g = ops.Graph()
   with g.as_default():
     features = {
         feature_keys.TrainEvalFeatures.TIMES:
             array_ops.zeros((1, 1)),
         feature_keys.TrainEvalFeatures.VALUES:
             array_ops.zeros((1, 1, 1)),
         "ticker":
             array_ops.reshape(
                 math_ops.cast(
                     variables.VariableV1(
                         name="ticker",
                         initial_value=0,
                         dtype=dtypes.int64,
                         collections=[ops.GraphKeys.LOCAL_VARIABLES])
                     .count_up_to(10),
                     dtype=dtypes.float32), (1, 1, 1))
     }
     model_fn = ts_head_lib.TimeSeriesRegressionHead(
         model=_TickerModel(),
         state_manager=state_management.PassthroughStateManager(),
         optimizer=train.GradientDescentOptimizer(0.001)).create_estimator_spec
     outputs = model_fn(
         features=features, labels=None, mode=estimator_lib.ModeKeys.EVAL)
     metric_update_ops = [
         metric[1] for metric in outputs.eval_metric_ops.values()]
     loss_mean, loss_update = metrics.mean(outputs.loss)
     metric_update_ops.append(loss_update)
     with self.cached_session() as sess:
       coordinator = coordinator_lib.Coordinator()
       queue_runner_impl.start_queue_runners(sess, coord=coordinator)
       variables.local_variables_initializer().run()
       sess.run(metric_update_ops)
       loss_evaled, metric_evaled, nested_metric_evaled = sess.run(
           (loss_mean, outputs.eval_metric_ops["ticker"][0],
            outputs.eval_metric_ops[feature_keys.FilteringResults.STATE_TUPLE][
                0][0]))
       # The custom model_utils metrics for in-sample predictions should be in
       # sync with the Estimator's mean metric for model loss.
       self.assertAllClose(0., loss_evaled)
       self.assertAllClose((((0.,),),), metric_evaled)
       self.assertAllClose((((0.,),),), nested_metric_evaled)
       coordinator.request_stop()
       coordinator.join()