code: def compute_retrieval_metrics(
      scores: jnp.ndarray,  # Shape: (N, M)
      target: Optional[jnp.ndarray] = None,  # Shape: (N,)
      k_values: Sequence[int] = (1, 5, 10),
      suffix: str = "",
      suffix_separator: str = "_",
  ) -> MutableMapping[str, float]:
      """Computes the text-to-visual retrieval metrics."""
      if target is None:
          assert scores.shape[0] == scores.shape[1], (
              "If the score matrix is not a"
              " square, the target values"
              " need to be provided as they"
              " can't be easily inferred."
          )
          target = jnp.arange(scores.shape[0])
  
      target = target.reshape(-1, 1)
  
      if suffix:
          suffix = suffix_separator + suffix
  
      summary = {}
  
      sorted_score_positions = scores.argsort(axis=-1)[..., ::-1]
      target_position_mask = sorted_score_positions == target
  
      for k in k_values:
          summary[f"recall@{k}{suffix}"] = (
              target_position_mask[:, :k].any(axis=-1).mean()
          )
  
      # We specify `size` so this function can be jitted.
      ranks = (
          jnp.argwhere(target_position_mask, size=len(scores))[:, -1]
          + 1
      )
      summary[f"ranks{suffix}"] = ranks
      summary[f"median_rank{suffix}"] = jnp.median(ranks)
      summary[f"mean_rank{suffix}"] = ranks.mean()
  
      return summary
  
imports: from collections.abc import Callable, Mapping, MutableMapping, Sequence
  import dataclasses
  from typing import Optional
  from clu import metric_writers
  from flax import jax_utils
  import flax.core
  import jax
  import jax.numpy
  import ml_collections
  from scenic.dataset_lib import dataset_utils
  from scenic.projects.lang4video.model import loss
  from scenic.projects.lang4video.model.image_text_model import ImageTextModel
  from scenic.projects.lang4video.trainer.train_utils import axis_name_exists
  from scenic.projects.lang4video.trainer.train_utils import compute_mask
  from scenic.projects.lang4video.trainer.train_utils import get_cached_fn
  from scenic.projects.lang4video.trainer.train_utils import get_epoch_steps
  from scenic.projects.lang4video.trainer.train_utils import get_input_spec
  from scenic.projects.lang4video.trainer.train_utils import init_encoder
  from scenic.projects.lang4video.trainer.train_utils import InputSpec
  from scenic.projects.lang4video.trainer.train_utils import is_video_input
  from scenic.projects.lang4video.trainer.train_utils import log_eval_summary
  from scenic.projects.lang4video.trainer.train_utils import NUM_DEVICES_AXIS_NAME
  from scenic.projects.lang4video.trainer.train_utils import pad_and_batch
  from scenic.projects.lang4video.trainer.train_utils import partial_with_cache
  from scenic.train_lib import train_utils
  from tqdm.auto import tqdm
label_desc: Use image node_id 2 for calling a global function with example usage: compute_retrieval_metrics(scores, target, k_values, suffix, suffix_separator) and returns: summary