class_docstr: A helper class to build the summary over the training iterations.
class_docstr_tok: ["A","helper","class","to","build","the","summary","over","the","training","iterations","."]
docstr: Write to the given writer and training step.
  
  After writing, the state gets reset.
  
  Args:
    writer: The summary will be written with this writer.
    step: The current training step.
  
  Returns:
    The summary since the last write.
docstr_tok: ["Write","to","the","given","writer","and","training","step",".","After","writing",",","the","state","gets","reset",".","Args",":","writer",":","The","summary","will","be","written","with","this","writer",".","step",":","The","current","training","step",".","Returns",":","The","summary","since","the","last","write","."]
code: def write(self, writer: metric_writers.MetricWriter, step: int):
      """Write to the given writer and training step.
  
      After writing, the state gets reset.
  
      Args:
        writer: The summary will be written with this writer.
        step: The current training step.
  
      Returns:
        The summary since the last write.
      """
      summary = train_utils.log_train_summary(
          step=step,
          train_metrics=jax.tree_util.tree_map(
              train_utils.unreplicate_and_get, self.metrics
          ),
          extra_training_logs=jax.tree_util.tree_map(
              train_utils.unreplicate_and_get, self.extra_logs
          ),
          writer=writer,
          key_separator="/",
      )
      self.metrics = []
      self.extra_logs = []
      return summary
  
imports: import copy
  import dataclasses
  import functools
  import json
  import os
  import random
  from typing import Any, Callable, Dict, List, Optional, Tuple
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from dmvr import tokenizers
  from flax import jax_utils
  from flax.core import unfreeze
  import flax.linen
  import jax
  import jax.example_libraries.optimizers
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import base_model
  from scenic.projects.t5 import model
  from scenic.projects.t5 import tokenizer
  from scenic.projects.vid2seq import load_utils
  from scenic.projects.vid2seq import models
  from scenic.projects.vid2seq import train_utils
  from scenic.projects.vid2seq import dvc_eval
  from scenic.train_lib_deprecated import lr_schedules
  from scenic.train_lib_deprecated import optimizers
  from scenic.train_lib_deprecated import pretrain_utils
  from scenic.train_lib_deprecated import train_utils
label_desc: Use image node_id 2 for calling the SummaryBuilder obj's underlying member method code with example usage: obj.write(writer, step) and returns: summary