Wrap a metric fn that returns scores and weights as an eval metric fn.

  The input metric_fn returns values for the current batch. The wrapper
  aggregates the return values collected over all of the batches evaluated.

  Args:
    metric_fn: function that returns scores and weights for the current batch's
      logits and predicted labels.

  Returns:
    function that aggregates the scores and weights from metric_fn.

def _convert_to_eval_metric(metric_fn):
  """Wrap a metric fn that returns scores and weights as an eval metric fn.

  The input metric_fn returns values for the current batch. The wrapper
  aggregates the return values collected over all of the batches evaluated.

  Args:
    metric_fn: function that returns scores and weights for the current batch's
      logits and predicted labels.

  Returns:
    function that aggregates the scores and weights from metric_fn.
  """
  def problem_metric_fn(*args):
    """Returns an aggregation of the metric_fn's returned values."""
    (scores, weights) = metric_fn(*args)

    # The tf.metrics.mean function assures correct aggregation.
    return tf.metrics.mean(scores, weights)
  return problem_metric_fn