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."""
  video_batches = [
      batch_samples['original_video'],
      batch_samples['generated_video'],
  ]
  if 'condition_video' in batch_samples:
    condition_video = jnp.where(
        batch_samples['condition_mask'], batch_samples['condition_video'],
        batch_samples['condition_video'] * 0.3)
    video_batches.append(condition_video)
  image_sample = train_utils.draw_frames_side_by_side(
      *video_batches, show_num=log_sample_size)
  image_samples = {f'{prefix}/generated_video_static': image_sample}
  writer.write_images(step, image_samples)
  if log_video:
    video_sample = train_utils.draw_videos_side_by_side(
        *video_batches, show_num=log_sample_size)
    video_samples = {f'{prefix}/generated_video': video_sample}
    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/maskgvt_trainer.py