def run_evals(self, summary_writers=None):
    """Runs and records evals for this training session.

    Args:
      summary_writers: List of per-task Jaxboard summary writers to log metrics.
    """
    if summary_writers is None:
      summary_writers = (None,) * len(self._eval_tasks)

    self._eval_model.weights = self._model.weights
    self._eval_model.state = self._model.state

    def recursively_look_for_printable_states(state):
      if isinstance(state, (tuple, list)):
        for substate in state:
          for item in recursively_look_for_printable_states(substate):
            yield item
      if isinstance(state, dict):
        for key, value in state.items():
          if isinstance(key, str) and key.startswith('summary_'):
            for device_id, device_value in enumerate(value):
              yield ('device{}/{}'.format(device_id, key[len('summary_'):]),
                     device_value)

    # The most recently trained weights are in this trainer, use those for eval.
    cur_train_task_index = self._which_task(self._step)
    trainer = self._trainer_per_task[cur_train_task_index]

    for eval_task_index in range(len(self._eval_tasks)):
      eval_task = self._eval_tasks[eval_task_index]
      evaluator = self._evaluator_per_task[eval_task_index]
      if eval_task is None:
        continue

      # Extract the actual model weights and state, excluding the loss layer.
      if self._use_memory_efficient_trainer:
        model_weights, model_state = self._model.weights, self._model.state
      else:
        model_weights = trainer.accelerated_model_with_loss.weights[0]
        model_state = trainer.accelerated_model_with_loss.state[0]

      # evaluator.{weights,state} are already replicated.
      metrics_weights = (model_weights, evaluator.weights)
      metrics_state = (model_state, evaluator.state)

      n_batches = eval_task.n_eval_batches
      n_metrics = len(eval_task.metrics)
      sums = np.zeros((n_metrics,))
      for _ in range(n_batches):
        rng = self.new_rng()
        batch = eval_task.next_batch()
        metric_values, _ = evaluator.metrics_fn(
            batch, metrics_weights, metrics_state, rng)
        sums += metric_values
      averages = sums / n_batches
      all_metrics = dict(zip(eval_task.metric_names, averages))
      summary_writer = summary_writers[eval_task_index]
      self.log_summary(all_metrics, summary_writer, 'metrics/', 'eval')
      summary_metrics = dict(recursively_look_for_printable_states(
          model_state))
      self.log_summary(summary_metrics, summary_writer, 'summary_', 'eval')
Runs and records evals for this training session.

    Args:
      summary_writers: List of per-task Jaxboard summary writers to log metrics.
    
trax/supervised/training.py