code: def multihead_classification_metrics_function(
      logits,
      batch,
      metrics: base_model.MetricNormalizerFnDict = _MULTIHEADLABEL_METRICS,
      class_splits: Optional[jnp.ndarray] = None,
      split_names: Optional[List[str]] = None,
  ) -> Dict[str, Any]:
      """Returns a callable metric function for the multihead classification task.
  
      Args:
        logits: Output of model in shape [batch, length, num_classes].
        batch: Batch of data that has 'label' and optionally 'batch_mask'.
        metrics: The multi-label classification metrics to evaluate. The key is the
          name of the  metric, and the value is the metrics function.
        class_splits: start indices of class splits.
        split_names: names of class splits.
  
      Returns:
        A dict of metrics, in which keys are metrics name and values are tuples of
        (metric, normalizer).
      """
  
      one_hot_targets = batch["label"]
      weights = batch.get(
          "batch_mask"
      )  # batch_mask might not be defined
  
      logit_splits = jnp.split(logits, class_splits, axis=-1)[:-1]
      one_hot_target_splits = jnp.split(
          one_hot_targets, class_splits, axis=-1
      )[:-1]
  
      evaluated_metrics = {}
      total_loss = [0.0, 0]
      for logits_i, one_hot_targets_i, name in zip(
          logit_splits, one_hot_target_splits, split_names
      ):
          for key, val in metrics.items():
              evaluated_metrics[
                  f"{name}_{key}"
              ] = model_utils.psum_metric_normalizer(  # pytype: disable=wrong-arg-types  # jax-ndarray
                  (
                      val[0](logits_i, one_hot_targets_i, weights),
                      val[1](logits_i, one_hot_targets_i, weights),
                  )
              )
              if key == "loss":
                  total_loss[0] += evaluated_metrics[f"{name}_{key}"][0]
                  total_loss[1] += evaluated_metrics[f"{name}_{key}"][1]
      evaluated_metrics["total_loss"] = tuple(total_loss)
  
      if len(class_splits) == 2:
          pairwise_acc = model_utils.psum_metric_normalizer(
              (
                  vivit_model_utils.joint_accuracy(
                      logits, one_hot_targets, class_splits, weights
                  ),
                  model_utils.num_examples(
                      logits, one_hot_targets, weights
                  ),
              )
          )
          pairwise_top_five = model_utils.psum_metric_normalizer(
              (
                  vivit_model_utils.joint_top_k(
                      logits,
                      one_hot_targets,
                      class_splits,
                      k=5,
                      weights=weights,
                  ),
                  model_utils.num_examples(
                      logits, one_hot_targets, weights
                  ),
              )
          )
          eval_name = f"{split_names[0]}-{split_names[1]}"
          evaluated_metrics[f"{eval_name}_accuracy"] = pairwise_acc
          evaluated_metrics[
              f"{eval_name}_accuracy_top_5"
          ] = pairwise_top_five
  
      return evaluated_metrics
  
imports: import functools
  from typing import Any, Dict, Optional, Tuple, List
  from absl import logging
  from flax.training import common_utils
  from immutabledict import immutabledict
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models.classification_model import classification_metrics_function
  from scenic.model_lib.base_models.encoder_decoder_model import encoder_decoder_metrics_function
  from scenic.model_lib.base_models.multilabel_classification_model import multilabel_classification_metrics_function
  from scenic.projects.vivit import model_utils
label_desc: Use image node_id 2 for calling a global function with example usage: multihead_classification_metrics_function(logits, batch, metrics, class_splits, split_names) and returns: evaluated_metrics