class_docstr: Plain Vision Transformer.
class_docstr_tok: ["Plain","Vision","Transformer","."]
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,
    batch)```
docstr_tok: ["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",",","batch",")","``","`"]
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,
        batch)```
      """
      del split  # For all splits, we return the same metric functions.
      loss_fn = self.config.get("loss", "sigmoid_xent")
      # pylint: disable=protected-access
      if loss_fn == "sigmoid_xent":
          return functools.partial(
              multilabel_classification_model.multilabel_classification_metrics_function,
              target_is_multihot=self.dataset_meta_data.get(
                  "target_is_onehot", False
              ),
              metrics=multilabel_classification_model._MULTI_LABEL_CLASSIFICATION_METRICS,
          )
      elif loss_fn == "softmax_xent":
          return functools.partial(
              classification_model.classification_metrics_function,
              target_is_onehot=self.dataset_meta_data.get(
                  "target_is_onehot", False
              ),
              metrics=classification_model._CLASSIFICATION_METRICS,
          )
      # pylint: enable=protected-access
      else:
          raise ValueError(f"Unknown loss function {loss_fn}.")
  
imports: import functools
  from typing import Any, Optional
  from absl import logging
  import flax
  import flax.linen
  from flax.training import common_utils
  import jax.numpy
  import ml_collections
  import numpy
  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 import multilabel_classification_model
  from scenic.model_lib.layers import nn_layers
  import scipy
label_desc: Use image node_id 3 for calling the PlainViT obj's underlying member method code with example usage: obj.get_metrics_fn(split) and returns: functools, functools