code: def get_tensorboard_log_fn(
      log_dir: str,
  ) -> Callable[[int, jnp.ndarray, str], None]:
      """
      Returns a custom logging function for the `evojax` `Trainer`.
      The function logs rewards after every train/test iteration with Tensorboard.
      It tries to use `tensorflow` or `pytorch` as tensorboard provider.
  
      Args:
          log_dir - directory to save store the tensorboard logs
      """
      try:
          from torch.utils.tensorboard import SummaryWriter
  
          def log_with_pytorch(i: int, scores: jnp.ndarray, stage: str):
              with SummaryWriter(log_dir=log_dir) as writer:
                  writer.add_scalar(
                      f"{stage}/score_min",
                      scores.min().item(),
                      global_step=i,
                  )
                  writer.add_scalar(
                      f"{stage}/score_max",
                      scores.max().item(),
                      global_step=i,
                  )
                  writer.add_scalar(
                      f"{stage}/score_mean",
                      scores.mean().item(),
                      global_step=i,
                  )
                  writer.add_scalar(
                      f"{stage}/score_std",
                      scores.std().item(),
                      global_step=i,
                  )
                  writer.add_histogram(
                      f"{stage}/score_distribution",
                      np.array(scores),
                      global_step=i,
                  )
  
          return log_with_pytorch
  
      except ImportError:
          pass
  
      try:
          import tensorflow as tf
  
          def log_with_tf(i: int, scores: jnp.ndarray, stage: str):
              with tf.summary.SummaryWriter(
                  log_dir=log_dir
              ).as_default():
                  tf.summary.scalar(
                      f"{stage}/score_min", scores.min().item(), step=i
                  )
                  tf.summary.scalar(
                      f"{stage}/score_max", scores.max().item(), step=i
                  )
                  tf.summary.scalar(
                      f"{stage}/score_mean",
                      scores.mean().item(),
                      step=i,
                  )
                  tf.summary.scalar(
                      f"{stage}/score_std", scores.std().item(), step=i
                  )
                  tf.summary.histogram(
                      f"{stage}/score_distribution",
                      np.array(scores),
                      step=i,
                  )
  
          return log_with_tf
  
      except ImportError:
          pass
  
      raise ImportError(
          "Please install the tensorboard AND (tensorflow OR pytorch) "
          "packages to log the rewards to tensorboard"
      )
  
file_path: evojax/evojax/util.py
imports: import os
  import logging
  import numpy
  from typing import Union
  from typing import Tuple
  from typing import Callable
  import jax.numpy
  from jax import tree_util
  from flax.core import FrozenDict
label_desc: Use image node_id 6 for calling a global function with example usage: get_tensorboard_log_fn(log_dir) and returns: log_with_pytorch, log_with_tf