code: def metrics_function_jit(
      logits: jnp.ndarray,
      batch: Batch,
      target_is_one_or_multihot: bool,
      metrics,
  ) -> Dict[str, Tuple[float, int]]:
      """Calculates metrics for the multi-label classification task for jit.
  
      Currently we assume each metric_fn has the API:
        ```metric_fn(logits, targets, weights)```
      and returns an array of shape [batch_size].
  
      Pmap-based trainers assume that to compute the aggregate metric, one should
      sum across all batches, then divide by the total samples seen.
  
      We follow the same API here, but note that summing should no longer use
      lax.psum, but rather a jnp.sum suffices as jit uses global arrays.
  
      Args:
       logits: Output of model in shape [batch, length, num_classes].
       batch: Batch of data that has 'label' and optionally 'batch_mask'.
       target_is_one_or_multihot: If the target is a one-hot or multi-hot vector.
       metrics: The 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_one_or_multihot:
          one_or_multihot_target = batch["label"]
      else:
          # This is to support running a multi-label classification model on
          # single-label classification tasks:
          one_or_multihot_target = common_utils.onehot(
              batch["label"], logits.shape[-1]
          )
      weights = batch.get(
          "batch_mask"
      )  # batch_mask might not be defined
  
      evaluated_metrics = {}
      for key, metric in metrics.items():
          fn, normaliser = metric
          metric_value = fn(logits, one_or_multihot_target, weights)
          norm_value = normaliser(
              logits, one_or_multihot_target, weights
          )
          evaluated_metrics[key] = (
              jnp.sum(metric_value),
              jnp.sum(norm_value),
          )
  
      return evaluated_metrics
  
file_path: scenic/scenic/model_lib/base_models/base_model.py
imports: from typing import Any, Callable, Dict, Mapping, Optional, Tuple
  from absl import logging
  import flax.linen
  from flax.training import common_utils
  import jax.numpy
  import ml_collections
label_desc: Use image node_id 1 for calling a global function with example usage: metrics_function_jit(logits, batch, target_is_one_or_multihot, metrics) and returns: evaluated_metrics