class_docstr: Video Transformer model for multiple n-way classification.
class_docstr_tok: ["Video","Transformer","model","for","multiple","n-way","classification","."]
docstr: Returns a callable metric function for the model.
  
  Args:
    split: The split for which we calculate the metrics. It should be one
      of the ['train',  'validation', 'test'].
  Returns: A metric function with the following API: ```metrics_fn(logits,
    label, weights)```
code: def get_metrics_fn(
      self, split: Optional[str] = None
  ) -> base_model.MetricFn:
      """Returns a callable metric function for the model.
  
      Args:
        split: The split for which we calculate the metrics. It should be one
          of the ['train',  'validation', 'test'].
      Returns: A metric function with the following API: ```metrics_fn(logits,
        label, weights)```
      """
      del split  # for all splits, we return the same metric functions
  
      def classification_metrics_function(
          logits, batch, metrics, class_splits, split_names
      ):
          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.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}"
                  ] = base_model_utils.psum_metric_normalizer(
                      (
                          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"] = total_loss
  
          if len(class_splits) == 2:
              pairwise_acc = base_model_utils.psum_metric_normalizer(
                  (
                      model_utils.joint_accuracy(
                          logits, one_hot_targets, class_splits, weights
                      ),
                      base_model_utils.num_examples(
                          logits, one_hot_targets, weights
                      ),
                  )
              )
              pairwise_top_five = (
                  base_model_utils.psum_metric_normalizer(
                      (
                          model_utils.joint_top_k(
                              logits,
                              one_hot_targets,
                              class_splits,
                              k=5,
                              weights=weights,
                          ),
                          base_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
  
      return functools.partial(
          classification_metrics_function,
          metrics=ViViT_CLASSIFICATION_METRICS,
          class_splits=self.class_splits,
          split_names=self.split_names,
      )
  
imports: import functools
  from typing import Any, Optional, Callable, Sequence
  from absl import logging
  import flax.linen
  from flax.linen.linear import default_kernel_init
  from immutabledict import immutabledict
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.common_lib import video_utils
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
  from scenic.projects.vivit import model_utils
label_desc: Use image node_id 3 for calling the ViViTMultiHeadClassificationModel obj's underlying member method code with example usage: obj.get_metrics_fn(split) and returns: functools, evaluated_metrics