code: def temporal_localization_metrics_function(
      logits: jnp.ndarray,
      batch: Batch,
      config: ml_collections.ConfigDict,
      classification_metrics: Mapping[
          str, Any
      ] = _TEMPORAL_LOCALIZATION_SIGMOID_LOSS_CLASSIFICATION_METRICS,
      box_regression_metrics: Mapping[
          str, Any
      ] = _TEMPORAL_LOCALIZATION_BOX_REGRESSION_METRICS,
      axis_name: Union[str, Tuple[str, ...]] = "batch",
  ) -> Dict[str, Tuple[float, float]]:
      """Calculates metrics for the temporal localization task.
  
      Args:
       logits: Output of model in shape [batch, num_frames, num_classes * 3] if
         config.output_per_class_displacements = True, otherwise in shape (batch,
         num_frames, num_classes + 2).
       batch: Batch of data that has 'label', 'displacements', and optionally
         'batch_mask'.
       config: Loss config.
       classification_metrics: Mapping from classification metric names to metric
         functions.
       box_regression_metrics: Mapping from box regression metric names to metric
         functions.
       axis_name: List of axes on which we run the pmsum.
  
      Returns:
        A dict of metrics, in which keys are metrics name and values are tuples of
        (metric, normalizer).
      """
      if batch.get("batch_mask") is None:
          batch_mask = jnp.ones((logits.shape[0],), dtype=jnp.float32)
      else:
          batch_mask = batch.get("batch_mask")
      weights = batch_mask[:, None] * batch["inputs"][
          "input_mask"
      ].astype(jnp.float32)
  
      # 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 = {}
      bs, num_frames, _ = logits.shape
      if config.get("output_per_class_displacements", True):
          num_classes = logits.shape[-1] // 3
          reshaped_logits = logits.reshape(
              (bs, num_frames, num_classes, 3)
          )
          class_logits = reshaped_logits[..., 0]
          pred_displacements = reshaped_logits[..., 1:]
          gt_displacements = batch["displacements"]
      else:
          num_classes = logits.shape[-1] - 2
          class_logits = logits[..., :num_classes]
          pred_displacements = jnp.expand_dims(
              logits[..., num_classes:], axis=2
          )
          gt_displacements = jnp.expand_dims(
              batch["displacements"], axis=2
          )
  
      class_label = batch["label"]
      for key, val in classification_metrics.items():
          if key == "focal_classification_loss":
              evaluated_metrics[
                  key
              ] = model_utils.psum_metric_normalizer(
                  (
                      val[0](
                          class_logits,
                          class_label,
                          weights,
                          alpha=config.get("focal_loss_alpha", 0.5),
                          gamma=config.get("focal_loss_gamma", 2.0),
                      ),
                      val[1](class_logits, class_label, weights),
                  ),
                  axis_name=axis_name,
              )
          else:
              evaluated_metrics[
                  key
              ] = model_utils.psum_metric_normalizer(
                  (
                      val[0](class_logits, class_label, weights),
                      val[1](class_logits, class_label, weights),
                  ),
                  axis_name=axis_name,
              )
  
      for key, val in box_regression_metrics.items():
          evaluated_metrics[key] = model_utils.psum_metric_normalizer(
              (
                  val[0](
                      pred_displacements,
                      gt_displacements,
                      class_label,
                      weights,
                  ),
                  val[1](class_label, weights),
              ),
              axis_name=axis_name,
          )
      return evaluated_metrics
  
imports: import functools
  from typing import Any, Callable, Dict, Mapping, Optional, Tuple, Union
  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.projects.unloc import metrics
label_desc: Use image node_id 7 for calling a global function with example usage: temporal_localization_metrics_function(logits, batch, config, classification_metrics, box_regression_metrics, axis_name) and returns: evaluated_metrics