examples/flax/image-captioning/run_image_captioning_flax.py/write_metric
def write_metric(summary_writer, metrics, train_time, step, metric_key_prefix="train"):
    if train_time:
        summary_writer.scalar("train_time", train_time, step)

        metrics = get_metrics(metrics)
        for key, vals in metrics.items():
            tag = f"{metric_key_prefix}_{key}"
            for i, val in enumerate(vals):
                summary_writer.scalar(tag, val, step - len(vals) + i + 1)

    else:
        for metric_name, value in metrics.items():
            summary_writer.scalar(f"{metric_key_prefix}_{metric_name}", value, step)