code: def polyvit_metrics_function(
      logits: jnp.ndarray,
      batch: base_model.Batch,
      dataset_name: str,
      dataset_meta_data: Dict[str, Any],
      class_splits: Dict[str, Any],
      split_names: Dict[str, Any],
  ) -> Dict[str, Tuple[float, int]]:
      """Defines and computes metrics for the polyvit model.
  
      Args:
       logits: Output of model in shape [batch, length, num_classes].
       batch: Batch of data that has 'label' and optionally 'batch_mask'.
       dataset_name: The name of the dataset used for the task.
       dataset_meta_data: Metadata of the dataset we are using.
       class_splits: start indices of class splits for multi-head classification
         tasks.
       split_names: names of class splits for multi-head classification tasks.
  
      Returns:
        A dict of metrics, in which keys are metrics name and values are tuples of
        (metric, normalizer).
      """
  
      current_dataset_meta_data = dataset_meta_data[dataset_name]
      task = current_dataset_meta_data["task"]
  
      kwargs = {}
  
      if task in [Task.MULTILABEL, Task.BOW]:
          kwargs = {
              "target_is_multihot": current_dataset_meta_data.get(
                  "target_is_onehot", False
              )
          }
      elif task == Task.MULTIHEADLABEL:
          kwargs = {
              "class_splits": class_splits[dataset_name],
              "split_names": split_names[dataset_name],
          }
      else:
          kwargs = {
              "target_is_onehot": current_dataset_meta_data.get(
                  "target_is_onehot", False
              )
          }
  
      return _METRICS_FUNCTIONS[task](logits, batch, **kwargs)
  
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 4 for calling a global function with example usage: polyvit_metrics_function(logits, batch, dataset_name, dataset_meta_data, class_splits, split_names) and returns: _METRICS_FUNCTIONS