def _open_summary_writers(self):
    """Opens the Jaxboard summary writers wrapped by context manager.

    Yields:
      A pair (train_summary_writers, eval_summary_writers) of lists of
      Jaxboard summary writers wrapped in a GeneratorContextManager object.
      Elements of the lists correspond to the training and evaluation task
      directories created during initialization. If there was no output_dir
      provided, yields lists of Nones with the appropriate length.
    """
    if self._output_dir is not None:
      _log(f'Metrics will be written in {self._output_dir}.', stdout=False)
      train_writers = [jaxboard.SummaryWriter(os.path.join(output_dir, 'train'))
                       for output_dir in self._output_dir_per_train_task]
      eval_writers = [jaxboard.SummaryWriter(os.path.join(output_dir, 'eval'))
                      for output_dir in self._output_dir_per_eval_task]
      try:
        yield (train_writers, eval_writers)
      finally:
        for writer in train_writers + eval_writers:
          writer.close()
        _log(f'Metrics were written in {self._output_dir}', stdout=False)
    else:
      yield ([None] * len(self._tasks), [None] * len(self._eval_tasks))
Opens the Jaxboard summary writers wrapped by context manager.

    Yields:
      A pair (train_summary_writers, eval_summary_writers) of lists of
      Jaxboard summary writers wrapped in a GeneratorContextManager object.
      Elements of the lists correspond to the training and evaluation task
      directories created during initialization. If there was no output_dir
      provided, yields lists of Nones with the appropriate length.
    
trax/supervised/training.py