code: def write_train_metric(
      summary_writer, train_metrics, train_time, step
  ):
      summary_writer.scalar("train_time", train_time, step)
  
      train_metrics = get_metrics(train_metrics)
      for key, vals in train_metrics.items():
          tag = f"train_{key}"
          for i, val in enumerate(vals):
              summary_writer.scalar(tag, val, step - len(vals) + i + 1)
  
imports: import logging
  import sys
  import time
  from dataclasses import field
  from pathlib import Path
  from typing import Dict, List, Optional, Union
  import numpy
  from datasets import DatasetDict, load_dataset
  from tqdm import tqdm
  import flax
  import jax
  import jax.numpy
  import librosa
  import optax
  from flax import jax_utils, traverse_util
  from flax.training import train_state
  from flax.training.common_utils import get_metrics, onehot, shard
  from transformers import FlaxWav2Vec2ForPreTraining, HfArgumentParser, TrainingArguments, Wav2Vec2Config, Wav2Vec2FeatureExtractor, is_tensorboard_available
  from transformers.models.wav2vec2.modeling_flax_wav2vec2 import _compute_mask_indices, _sample_negative_indices
label_desc: Use image node_id 2 for calling a global function with example usage: write_train_metric(summary_writer, train_metrics, train_time, step) without return types