code: def global_metrics_fn(
      all_confusion_mats: List[jnp.ndarray],
      dataset_metadata: Dict[str, Any],
  ) -> Dict[str, float]:
      """Returns a dict with global (whole-dataset) metrics."""
      # Compute mIoU from list of confusion matrices:
      assert isinstance(
          all_confusion_mats, list
      )  # List of eval batches.
      cm = np.sum(all_confusion_mats, axis=0)  # Sum over eval batches.
      assert cm.ndim == 3, (
          "Expecting confusion matrix to have shape "
          "[batch_size, num_classes, num_classes], got "
          f"{cm.shape}."
      )
      cm = np.sum(cm, axis=0)  # Sum over batch dimension.
      mean_iou, iou_per_class = model_utils.mean_iou(cm)
      metrics_dict = {"mean_iou": float(mean_iou)}
      for label, iou in enumerate(iou_per_class):
          tag = f"iou_per_class/{label:02.0f}"
          if "class_names" in dataset_metadata:
              tag = f"{tag}_{dataset_metadata['class_names'][label]}"
          metrics_dict[tag] = float(iou)
      return metrics_dict
  
file_path: scenic/scenic/model_lib/base_models/segmentation_model.py
imports: import functools
  from typing import Any, Callable, List, Dict, Optional, Tuple, Union
  from flax.training import common_utils
  from immutabledict import immutabledict
  import jax.numpy
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
label_desc: Use image node_id 3 for calling a global function with example usage: global_metrics_fn(all_confusion_mats, dataset_metadata) and returns: metrics_dict