Create the metrics_fn that TPUEstimatorSpec expects.

def create_tpu_eval_metrics_fn(problem, model_hparams):
  """Create the metrics_fn that TPUEstimatorSpec expects."""

  metric_fns = []
  eval_metrics = problem.eval_metric_fns(model_hparams)

  tm = _create_target_modality(problem.get_hparams(model_hparams).modality)
  if isinstance(tm, dict):
    for k, v in six.iteritems(tm):
      weights_fn = modalities.get_weights_fn(v)

      def make_metric_fn(metric_fn):
        def wrapped_metric_fn(logits, labels, features, weights_fn=weights_fn):
          kwargs = {}
          args, _, keywords, _ = inspect.getargspec(metric_fn)
          if ("features" in args) or keywords:
            kwargs["features"] = features
          num, den = metric_fn(logits, labels, weights_fn=weights_fn, **kwargs)
          return tf.metrics.mean(num, den)

        return wrapped_metric_fn

      for metric, metric_fn in six.iteritems(eval_metrics):
        if metric in TPU_METRIC_BLACKLIST:
          log_warn("Skipping eval metric %s in TPU_METRIC_BLACKLIST", metric)
          continue
        name = "%s/metrics-%s/%s" % (k, problem.name, metric)
        metric_fns.append((name, make_metric_fn(metric_fn)))
  else:
    weights_fn = modalities.get_weights_fn(tm)

    def make_metric_fn(metric_fn):
      def wrapped_metric_fn(logits, labels, features):
        kwargs = {}
        args, _, keywords, _ = inspect.getargspec(metric_fn)
        if ("features" in args) or keywords:
          kwargs["features"] = features
        num, den = metric_fn(logits, labels, weights_fn=weights_fn, **kwargs)
        return tf.metrics.mean(num, den)

      return wrapped_metric_fn

    for metric, metric_fn in six.iteritems(eval_metrics):
      if metric in TPU_METRIC_BLACKLIST:
        log_warn("Skipping eval metric %s in TPU_METRIC_BLACKLIST", metric)
        continue
      name = "metrics-%s/%s" % (problem.name, metric)
      metric_fns.append((name, make_metric_fn(metric_fn)))

  def all_metrics_fn(**kwargs):
    """Construct metrics dictionary."""

    original_kwargs = _unflatten_dict(kwargs, prefixes=["logits", "features"])
    del kwargs

    logits = original_kwargs["logits"]
    labels = original_kwargs["labels"]
    features = original_kwargs["features"]
    del original_kwargs

    metrics_dict = {}

    for name, fn in metric_fns:
      if isinstance(logits, dict) and isinstance(labels, dict):
        for k, v in six.iteritems(logits):
          metrics_dict["%s/%s" % (k, name)] = fn(v, labels[k], features)
      elif isinstance(logits, dict):
        tf.logging.warning("Logits is a dict, but labels is not; only "
                           "evaluating logits['targets'] against labels.")
        metrics_dict["%s/%s" % ("targets", name)] = fn(logits["targets"],
                                                       labels, features)
      else:
        metrics_dict[name] = fn(logits, labels, features)

    return metrics_dict

  return all_metrics_fn