code: def _compute_metrics(
      encoded_visuals: jnp.ndarray,  # Shape: (N, E)
      encoded_texts: jnp.ndarray,  # Shape: (N, E)
      model: ImageTextModel,
      loss_fn: Callable[..., jnp.ndarray],
      retrieval_batch_size: Optional[int] = None,
  ) -> MutableMapping[str, float]:
      """Computes the similarity and the metrics."""
      assert encoded_visuals.shape == encoded_texts.shape
  
      n = encoded_texts.shape[0]
  
      # We don't batch because of potential OOMs but because we may want to compute
      # the metrics with a different number of negatives.
      retrieval_batch_size = retrieval_batch_size or n
  
      # We don't use `batch_and_vmap` here because we need to use the batched output
      # with some special padding values for the loss.
      batched, batch_padding_size = pad_and_batch(
          (encoded_texts, encoded_visuals),
          batch_size=retrieval_batch_size,
      )
      batched_scores = jax.vmap(model.flax_model.compute_similarity)(
          *batched
      )
  
      n_batches = batched_scores.shape[0]
  
      # We set the padding columns to -inf, so we can ignore them in the scores.
      batched_scores = batched_scores.at[
          -1, :, retrieval_batch_size - batch_padding_size :
      ].set(-jnp.inf)
  
      scores = batched_scores.reshape(-1, retrieval_batch_size)[:n]
      target = jnp.tile(jnp.arange(retrieval_batch_size), n_batches)[:n]
      metrics_summary = compute_retrieval_metrics(scores, target)
  
      where = jnp.ones((n_batches, retrieval_batch_size)).astype(
          jnp.bool_
      )
      where = where.at[
          -1, retrieval_batch_size - batch_padding_size :
      ].set(False)
  
      initial = jnp.full((n_batches,), -jnp.inf)
      metrics_summary["loss"] = jax.vmap(loss_fn)(
          batched_scores, where=where, initial=initial
      ).mean(where=where)
  
      return metrics_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 3 for calling a global function with example usage: _compute_metrics(encoded_visuals, encoded_texts, model, loss_fn, retrieval_batch_size) and returns: metrics_summary