code: def evaluate(inputs_stream, predict_fun, metric_funs, rng):
    
    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)}
docs: 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.