Evaluate.

  Args:
    inputs_stream: iterable of inputs to evaluate on.
    predict_fun: function from inputs to predictions. params should already be
      partially applied.
    metric_funs: dict from metric name to metric function, which takes inputs
      and predictions and returns a scalar metric value.
    rng: random number generator.

  Returns:
    metrics: dict from metric name to metric value averaged over the number of
      inputs.

def evaluate(inputs_stream, predict_fun, metric_funs, rng):
  """Evaluate.

  Args:
    inputs_stream: iterable of inputs to evaluate on.
    predict_fun: function from inputs to predictions. params should already be
      partially applied.
    metric_funs: dict from metric name to metric function, which takes inputs
      and predictions and returns a scalar metric value.
    rng: random number generator.

  Returns:
    metrics: dict from metric name to metric value averaged over the number of
      inputs.
  """
  metrics = collections.defaultdict(float)
  count = 0
  for inp in inputs_stream:
    count += 1
    rng, subrng = jax_random.split(rng)
    preds = predict_fun(inp[0], rng=subrng)
    for m, f in six.iteritems(metric_funs):
      metrics[m] += f(inp, preds)
  return {m: v / count for (m, v) in six.iteritems(metrics)}