code: def _log_to_tensorboard(
      writer: tensorboard.SummaryWriter,
      state: model_utils.TrainState,
      scalar_params: training.ScalarParams,
      stats: Dict[str, Union[Dict[str, jnp.ndarray], jnp.ndarray]],
      time_dict: Dict[str, jnp.ndarray],
  ):
      """Log statistics to Tensorboard."""
      step = int(state.optimizer.state.step)
      writer.scalar(
          "params/learning_rate", scalar_params.learning_rate, step
      )
      writer.scalar("params/warp_alpha", state.warp_alpha, step)
      writer.scalar("params/time_alpha", state.time_alpha, step)
      writer.scalar(
          "params/elastic_loss/weight",
          scalar_params.elastic_loss_weight,
          step,
      )
  
      # pmean is applied in train_step so just take the item.
      for branch in {"coarse", "fine"}:
          if branch not in stats:
              continue
          for stat_key, stat_value in stats[branch].items():
              writer.scalar(f"{stat_key}/{branch}", stat_value, step)
  
      if "background_loss" in stats:
          writer.scalar(
              "loss/background", stats["background_loss"], step
          )
  
      for k, v in time_dict.items():
          writer.scalar(f"time/{k}", v, step)
  
imports: import functools
  from typing import Dict, Union
  from absl import app
  from absl import flags
  from absl import logging
  from flax import jax_utils
  from flax import optim
  from flax.metrics import tensorboard
  from flax.training import checkpoints
  import gin
  import jax
  from jax import numpy
  from jax import random
  import numpy
  import tensorflow
  from nerfies import configs
  from nerfies import datasets
  from nerfies import gpath
  from nerfies import model_utils
  from nerfies import models
  from nerfies import schedules
  from nerfies import training
  from nerfies import utils
label_desc: Use image node_id 1 for calling a global function with example usage: _log_to_tensorboard(writer, state, scalar_params, stats, time_dict) without return types