code: def calculate_iou(predictions, labels, n_classes):
      """Calculates mean IoU of the entire test set."""
      all_intersection = np.zeros(n_classes)
      all_union = np.zeros(n_classes)
      for sem_idx in range(labels.shape[0]):
          for sem in range(n_classes):
              intersection = np.sum(
                  np.logical_and(
                      predictions[sem_idx] == sem,
                      labels[sem_idx] == sem,
                  )
              )
              union = jnp.sum(
                  np.logical_or(
                      predictions[sem_idx] == sem,
                      labels[sem_idx] == sem,
                  )
              )
              all_intersection[sem] += intersection
              all_union[sem] += union
      return np.mean(all_intersection / all_union)
  
imports: import functools
  from typing import Any, Callable, Dict, Optional, Tuple, Type
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from flax import jax_utils
  import flax.linen
  import jax
  from jax.example_libraries.optimizers import clip_grads
  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.model_lib.base_models import model_utils
  from scenic.projects.vivit import evaluation_lib
  from scenic.train_lib_deprecated import lr_schedules
  from scenic.train_lib_deprecated import optimizers
  from scenic.train_lib_deprecated import train_utils
label_desc: Use image node_id 4 for calling a global function with example usage: calculate_iou(predictions, labels, n_classes) and returns: np