def log_samples(batch_samples: Batch, writer: metric_writers.MetricWriter,
                prefix: str, step: int, log_sample_size: int,
                log_video: bool = True):
  """Logs a batch of samples as static image and video."""
  image_samples = {}
  for suffix in ['', '_ema']:
    key = f'generated_video{suffix}'
    image_samples[key] = train_utils.draw_frames_side_by_side(
        batch_samples['original_video'],
        batch_samples[key],
        show_num=log_sample_size)
  image_samples = {
      f'{prefix}/{k}_static': v for k, v in image_samples.items()
  }
  writer.write_images(step, image_samples)
  if log_video:
    video_samples = {}
    for suffix in ['', '_ema']:
      key = f'generated_video{suffix}'
      video_samples[key] = train_utils.draw_videos_side_by_side(
          batch_samples['original_video'],
          batch_samples[key],
          show_num=log_sample_size)
    video_samples = {f'{prefix}/{k}': v for k, v in video_samples.items()}
    video_samples = jax.tree_util.tree_map(
        lambda x: (x * 255.).astype(jnp.uint8), video_samples)
    writer.write_videos(step, video_samples)
Logs a batch of samples as static image and video.
videogvt/trainers/vqgan_trainer.py