code: def encoder_decoder_metrics_function(
      logits: jnp.ndarray,
      batch: base_model.Batch,
      target_is_onehot: bool = False,
      metrics: base_model.MetricNormalizerFnDict = _ENCODER_DECODER_METRICS,
      axis_name: Union[str, Tuple[str, ...]] = "batch",
  ) -> Dict[str, Tuple[float, int]]:
      """Calculates metrics for the encoder-decoder models.
  
  
      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].
       batch: Batch of data that has 'label' and optionally 'batch_mask'.
       target_is_onehot: If the target is a one-hot vector.
       metrics: The encoder-decoder metrics to evaluate. The key is the name of the
         metric, and the value is the metrics function.
       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 target_is_onehot:
          one_hot_targets = batch["label"]
      else:
          one_hot_targets = common_utils.onehot(
              batch["label"], logits.shape[-1]
          )
      weights = batch.get(
          "batch_mask"
      )  # `batch_mask` might not be defined.
  
      # Expanding from sequence-level to token level masking.
      if weights is not None:
          weights = jnp.tile(
              jnp.expand_dims(weights, axis=1), one_hot_targets.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_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 = {}
      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, weights),
                  val[
                      1
                  ](  # pytype: disable=wrong-arg-types  # jax-types
                      logits, one_hot_targets, weights
                  ),
              ),
              axis_name=axis_name,
          )
          if key == "loss":
              # TODO(dehghani): Move this to the training loop.
              # Calculate (clipped) perplexity after averaging log-perplexities:
              evaluated_metrics["perplexity"] = (
                  jnp.clip(
                      jnp.exp(
                          evaluated_metrics["loss"][0]
                          / evaluated_metrics["loss"][1]
                      ),
                      a_max=_MAX_PERPLEXITY,
                  ),
                  1,
              )
      return evaluated_metrics
  
file_path: scenic/scenic/model_lib/base_models/encoder_decoder_model.py
imports: import functools
  from typing import Dict, Optional, Tuple, Union
  from flax.training import common_utils
  from immutabledict import immutabledict
  import jax.numpy
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
label_desc: Use image node_id 2 for calling a global function with example usage: encoder_decoder_metrics_function(logits, batch, target_is_onehot, metrics, axis_name) and returns: evaluated_metrics