class_docstr: Sequence to sequence model.
class_docstr_tok: ["Sequence","to","sequence","model","."]
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,
    targets, weights)```
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",",","targets",",","weights",")","``","`"]
code: def get_metrics_fn(self, split: Optional[str] = None):
      """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,
        targets, weights)```
      """
      del split  # The metric function is the same for all splits.
  
      def metric_fn(
          logits: jnp.ndarray,
          targets: jnp.ndarray,
          weights: jnp.ndarray,
          target_is_onehot: bool = False,
          metrics: base_model.MetricNormalizerFnDict = _CLASSIFICATION_METRICS,
      ) -> Dict[str, Tuple[float, int]]:
          """Calcualte metrics for the 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].
           targets: Targets to be decoded.
           weights: Indicate which tokens are valid (1) vs padding (0).
           target_is_onehot: If the target is a one-hot vector.
           metrics: The 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_onehot:
              one_hot_targets = targets
          else:
              one_hot_targets = common_utils.onehot(
                  targets, logits.shape[-1]
              )
  
          # 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_devices, batch_size] -> [batch_size]
          # by summing across the devices dim. 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,
                          one_hot_targets,  # pytype: disable=wrong-arg-types  # jax-types
                          weights,
                      ),
                      val[1](logits, one_hot_targets, weights),
                  )
              )  # pytype: disable=wrong-arg-types  # jax-types
          return evaluated_metrics  # pytype: disable=bad-return-type  # jax-types
  
      return metric_fn
  
imports: from typing import Any, Dict, Optional, Tuple
  from absl import logging
  from flax import linen
  from flax.training import common_utils
  from immutabledict import immutabledict
  import jax
  from jax import lax
  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.mbt import model
  from scenic.projects.mbt.model import temporal_encode
label_desc: Use image node_id 3 for calling the Seq2SeqModel obj's underlying member method code with example usage: obj.get_metrics_fn(split) and returns: metric_fn, evaluated_metrics