code: def log_train_summary(
      step: int,
      *,
      writer: metric_writers.MetricWriter,
      train_metrics: Sequence[Dict[str, Tuple[float, int]]],
      extra_training_logs: Optional[Sequence[Dict[str, Any]]] = None,
      metrics_normalizer_fn: Optional[
          Callable[
              [Dict[str, Tuple[float, int]], str], Dict[str, float]
          ]
      ] = None,
      prefix: str = "train",
      key_separator: str = "_",
      flush_writer: bool = True,
  ) -> Dict[str, float]:
      """Computes and logs train metrics.
  
      Args:
        step: Current step.
        writer: Summary writer.
        train_metrics: List of dictionaries of calculated metrics. Usually the
          sequence is the concatenation of the per-eval-step metrics, and every
          dictionary maps a metric name to an array of (value, normalizer) - where
          the array index is usually the batch index.
        extra_training_logs: List of dictionaries, containing additional training
          logs, from every train step, e.g. learning rate, Time, num parameters,
          etc. Their mean will be logged.
        metrics_normalizer_fn: Used for normalizing metrics. The API for this
          function is: `new_metrics_dict = metrics_normalizer_fn(metrics_dict,
          split)`. If set to None, we use the normalize_metrics_summary which uses
          the normalizer paired with each metric to normalize it.
        prefix: str; Prefix added to the name of the summaries writen by this
          function.
        key_separator: Separator added between the prefix and key.
        flush_writer: If True, flush the writer after logging.
  
      Returns:
        A dictionary of metrics, mapping `train_metrics from metric name (incl.
        `prefix`) to float value.
      """
      ##### Prepare metrics:
      # Get metrics from devices:
      train_metrics = stack_forest(train_metrics)
      # Compute the sum over all examples in all batches:
      train_metrics_summary = jax.tree_util.tree_map(
          lambda x: x.sum(), train_metrics
      )
      # Normalize metrics by the total number of examples:
      metrics_normalizer_fn = (
          metrics_normalizer_fn or normalize_metrics_summary
      )
      train_metrics_summary = metrics_normalizer_fn(
          train_metrics_summary, "train"
      )
  
      ##### Prepare additional training logs:
      # If None, set to an empty dictionary.
      extra_training_logs = extra_training_logs or [{}]
      train_logs = stack_forest(extra_training_logs)
  
      # Metrics:
      writer.write_scalars(
          step,
          {
              key_separator.join((prefix, key)): val
              for key, val in train_metrics_summary.items()
          },
      )
      # Additional logs:
      writer.write_scalars(
          step, {key: val.mean() for key, val in train_logs.items()}
      )
  
      if flush_writer:
          writer.flush()
      return train_metrics_summary
  
imports: import collections.abc
  import copy
  import functools
  import os
  import re
  import time
  from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple, Union
  from absl import logging
  from clu import metric_writers
  import flax
  from flax import jax_utils
  from flax import struct
  import flax.linen
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  import optax
  from scenic.common_lib import debug_utils
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  from scenic.train_lib import optimizers
  from tensorflow.io import gfile
label_desc: Use image node_id 19 for calling a global function with example usage: log_train_summary(step) and returns: train_metrics_summary