code: def feature_regression_metrics_function(
      predictions: ArrayDict,
      prediction_masks: ArrayDict,
      batch: base_model_lib.Batch,
      feature_target: str,
      metrics: base_model_lib.MetricNormalizerFnDict,
  ) -> Dict[str, Tuple[float, int]]:
      """Calculate metrics for the feature 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:  Dictionary with the output of model in shape [batch, length].
        prediction_masks:  Dictionary of the predictions which are valid.
        batch: Batch (dict) with keys 'targets' and optionally 'batch_mask'.
        feature_target: The feature targets used for feature regression.
        metrics: The regression 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).
      """
      evaluated_metrics = {}
      # TODO(lgeorgescu): add the weighted sum in the evaluation too
  
      for key_target in feature_target:
          batch_mask = batch.get("batch_mask")
          if batch_mask is None:
              batch_mask = jnp.ones(prediction_masks[key_target].shape)
          if batch_mask.ndim == 1:
              n_batch = predictions[key_target].shape[0]
              batch_mask = jnp.reshape(batch_mask, (n_batch, 1))
          weights = batch_mask * prediction_masks[key_target]
  
          for key, val in metrics.items():
              evaluated_metrics[
                  key + "_" + key_target
              ] = model_utils_lib.psum_metric_normalizer(
                  (
                      val[0](
                          predictions[key_target],
                          batch["targets"][
                              key_target
                          ],  # pytype: disable=wrong-arg-types  # jax-ndarray
                          weights,
                      ),
                      val[1](
                          predictions[key_target],
                          batch["targets"][
                              key_target
                          ],  # pytype: disable=wrong-arg-types  # jax-ndarray
                          weights,
                      ),
                  )
              )
  
      return evaluated_metrics
  
imports: import functools
  from typing import Any, Dict, Optional, Tuple
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.av_mae import base_model
  from scenic.projects.av_mae import model_utils
  from scenic.projects.mbt import model
  from scenic.projects.vivit import model
  from scenic.projects.vivit import model_utils
label_desc: Use image node_id 5 for calling a global function with example usage: feature_regression_metrics_function(predictions, prediction_masks, batch, feature_target, metrics) and returns: evaluated_metrics