def log_summary(self, values, summary_writer, value_prefix, log_prefix,
                  stdout=True):
    """Logs and saves provided metrics.

    Args:
      values: Dict from metric name to metric value.
      summary_writer: Jaxboard summary writer.
      value_prefix: String appended in front of summary_writer entries.
      log_prefix: String appended in front of logs.
      stdout: Boolean saying if logs should be logged to stdout as well.
    """
    history = self._history
    should_write_summaries = self.is_chief and summary_writer is not None
    for name, value in values.items():
      full_name = value_prefix + name
      s = tuple(jnp.shape(value))
      if not s:
        self._log_step(
            '%s %s | % .8f' %
            (log_prefix.ljust(5), name.rjust(self._rjust_len), value),
            stdout=stdout)
        if should_write_summaries:
          summary_writer.scalar(full_name, value, self.step)
      else:
        if should_write_summaries:
          summary_writer.image(full_name, value, self.step)
      if history:
        history.append(log_prefix, full_name, self.step, value)
    if should_write_summaries:
      summary_writer.flush()
Logs and saves provided metrics.

    Args:
      values: Dict from metric name to metric value.
      summary_writer: Jaxboard summary writer.
      value_prefix: String appended in front of summary_writer entries.
      log_prefix: String appended in front of logs.
      stdout: Boolean saying if logs should be logged to stdout as well.
    
trax/supervised/training.py