def evaluation_round(self, inputs_stream, weights, state, rng):
    """Evaluate.

    Args:
      inputs_stream: Iterable of inputs to evaluate on.
      weights: Weights for each f in eval_fns.
      state: State for each f in eval_fns.
      rng: Single-use random number generator (JAX PRNG key).

    Returns:
      Tuple of `(metrics, state)`. `metrics` is a dict from metric name to
      metric value averaged over the number of inputs, and `state` is the end
      state returned by this trainer's `predict_fn`.
    """
    metrics = collections.defaultdict(float)
    count = 0
    for inp in inputs_stream:
      count += 1
      rng, subrng = jax_random.split(rng)
      metric_values, _ = self._jit_eval(inp, weights, state, subrng)
      try:
        metric_values = list(metric_values)
      except (TypeError, IndexError):
        metric_values = [float(metric_values)]
      for m, v in zip(self._metrics, metric_values):
        metrics[m] += v
    return {m: v / count for (m, v) in metrics.items()}, state
Evaluate.

    Args:
      inputs_stream: Iterable of inputs to evaluate on.
      weights: Weights for each f in eval_fns.
      state: State for each f in eval_fns.
      rng: Single-use random number generator (JAX PRNG key).

    Returns:
      Tuple of `(metrics, state)`. `metrics` is a dict from metric name to
      metric value averaged over the number of inputs, and `state` is the end
      state returned by this trainer's `predict_fn`.
    
trax/supervised/trainer_lib.py