code: def regression_metrics_function(
      predictions: jnp.ndarray,
      batch: base_model.Batch,
      metrics: base_model.MetricNormalizerFnDict = _REGRESSION_METRICS,
      axis_name: Union[str, Tuple[str, ...]] = "batch",
  ) -> Dict[str, Tuple[float, int]]:
      """Calculate metrics for the regression task.
  
      Currently we assume each metric_fn has the API:
        ```metric_fn(predictions, targets, weights)```
      and returns an array of shape [batch,]. 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:
       predictions: Output of model in shape [batch, length].
       batch: Batch (dict) with keys 'targets' and optionally 'batch_mask'.
       metrics: The regression 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).
      """
      targets = batch["targets"]
      weights = batch.get("batch_mask")
      evaluated_metrics = {}
      for key, val in metrics.items():
          evaluated_metrics[key] = model_utils.psum_metric_normalizer(
              (
                  val[0](predictions, targets, weights),
                  val[1](
                      predictions,
                      targets,  # pytype: disable=wrong-arg-types  # jax-ndarray
                      weights,
                  ),
              ),
              axis_name=axis_name,
          )
      return evaluated_metrics
  
file_path: scenic/scenic/model_lib/base_models/regression_model.py
imports: import functools
  from typing import Dict, Optional, Tuple, Union
  from immutabledict import immutabledict
  import jax.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 1 for calling a global function with example usage: regression_metrics_function(predictions, batch, metrics, axis_name) and returns: evaluated_metrics