code: def bow_classification_metrics_function(
      logits: jnp.ndarray,
      batch: base_model.Batch,
      target_is_multihot: bool = False,
      metrics: base_model.MetricNormalizerFnDict = _BOW_CLASSIFICATION_METRICS,
  ) -> Dict[str, Tuple[float, int]]:
      """Calcualte metrics for the Bag of Words classification task.
  
      Currently we assume each metric_fn has the API:
        ```metric_fn(logits, targets, weights)```
      and returns an array of shape [batch_size]. We also assume that to compute
      the aggregate metric, one should sum across all batches, then divide by the
      total samples seen. In this way we currently only support metrics of the 1/N
      sum f(inputs, targets). Note, the caller is responsible for dividing by
      the normalizer when computing the mean of each metric.
  
      Args:
       logits: Output of model in shape [batch, length, num_classes].
       batch: Batch of data that has 'label' and optionally 'batch_mask'.
       target_is_multihot: If the target is a multi-hot vector.
       metrics: The multi-label classification metrics to evaluate. The key is the
         name of the  metric, and the value is the metrics function.
  
      Returns:
        A dict of metrics, in which keys are metrics name and values are tuples of
        (metric, normalizer).
      """
      if target_is_multihot:
          multihot_target = batch["label"]
      else:
          # This is to support running a multi-label classification model on
          # single-label classification tasks:
          multihot_target = common_utils.onehot(
              batch["label"], logits.shape[-1]
          )
  
      # multihot_target is initially one-hot of shape (bs, len, vocab_size),
      # while we actually need multi-hot of shape (bs, vocab_size).
      multihot_target = multihot_target.max(axis=-2)
  
      weights = batch.get(
          "batch_mask"
      )  # batch_mask might not be defined
  
      # This psum is required to correctly evaluate with multihost. Only host 0
      # will report the metrics, so we must aggregate across all hosts. The psum
      # will map an array of shape [n_global_devices, batch_size] -> [batch_size]
      # by summing across the devices dimension. The outer sum then sums across the
      # batch dim. The result is then we have summed across all samples in the
      # sharded batch.
      evaluated_metrics = {}
      for key, val in metrics.items():
          evaluated_metrics[
              key
          ] = model_utils.psum_metric_normalizer(  # pytype: disable=wrong-arg-types  # jax-ndarray
              (
                  val[0](
                      logits, multihot_target, weights
                  ),  # pytype: disable=wrong-arg-types  # jax-types
                  val[1](logits, multihot_target, weights),
              )
          )  # pytype: disable=wrong-arg-types  # jax-types
      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 1 for calling a global function with example usage: bow_classification_metrics_function(logits, batch, target_is_multihot, metrics) and returns: evaluated_metrics