code: def eval_step(
      train_state: train_utils.TrainState,
      visual: jnp.ndarray,  # Shape: (N, F, H, W, C) if `is_video`, w/o F if not.
      text: jnp.ndarray,  # Shape: (N, L)
      mask: Optional[jnp.ndarray],  # Shape: (N, L)
      batch_mask: jnp.ndarray,
      text_for_mlm: jnp.ndarray,
      segment_ids_for_mlm: jnp.ndarray,
      mask_for_mlm: Optional[jnp.ndarray],
      masked_lm_positions: jnp.ndarray,
      masked_lm_ids: jnp.ndarray,
      masked_lm_weights: jnp.ndarray,
      text_batch_mask: jnp.ndarray,
      *,
      model: ImageTextModel,
      is_video: bool,
      gather_scores: bool = True,
      debug: bool = False,
  ) -> Mapping[str, Any]:
      """Runs a single step of evaluation."""
      encoder = model.flax_model
  
      encoded_visual, encoded_text, mlm_logits = encoder.apply(
          variables={
              "params": train_state.params,
              **train_state.model_state,
          },
          method=(
              model.flax_model.encode_video_and_text_and_do_pretraining
              if is_video
              else model.flax_model.encode_image_and_text_and_do_pretraining
          ),
          **{"video" if is_video else "image": visual},
          text=text,
          mask=mask,
          text_for_mlm=text_for_mlm,
          segment_ids_for_mlm=segment_ids_for_mlm,
          mask_for_mlm=mask_for_mlm,
          masked_lm_positions=masked_lm_positions,
          train=False,
          debug=debug,
      )
  
      scores = encoder.compute_similarity(
          encoded_text,
          encoded_visual,
          all_gather_axis_name=(
              NUM_DEVICES_AXIS_NAME
              if gather_scores
              and axis_name_exists(NUM_DEVICES_AXIS_NAME)
              else None
          ),
      )
  
      if gather_scores and axis_name_exists(NUM_DEVICES_AXIS_NAME):
          batch_mask = jax.lax.all_gather(
              batch_mask, NUM_DEVICES_AXIS_NAME
          )
          batch_mask = batch_mask.reshape(-1, batch_mask.shape[-1])
  
      loss_image_text = model.loss_function(scores)
      loss_mlm = bert_base_model.sparse_weighted_softmax_cross_entropy(
          mlm_logits, masked_lm_ids, masked_lm_weights, text_batch_mask
      )
      loss = loss_image_text + loss_mlm
  
      metrics = {
          "loss": loss,
          "loss/image_text": loss_image_text,
      }
  
      batch_mask_sum = batch_mask.sum()
      for k, v in metrics.items():
          metrics[k] = (v * batch_mask, batch_mask_sum)
  
      if not gather_scores and axis_name_exists(NUM_DEVICES_AXIS_NAME):
          # We need to gather the metrics then.
          actual_batch_size = jax.lax.psum(
              batch_mask_sum, axis_name=NUM_DEVICES_AXIS_NAME
          )
          for k, v in metrics.items():
              metrics[k] = (
                  jax.lax.psum(
                      v[0].sum(), axis_name=NUM_DEVICES_AXIS_NAME
                  ),
                  actual_batch_size,
              )
  
      metrics["loss/mlm"] = (
          loss_mlm * text_batch_mask,
          text_batch_mask.sum(),
      )
  
      return metrics
  
imports: from collections.abc import Callable
  from collections.abc import Mapping
  from collections.abc import MutableMapping
  import functools
  from typing import Any
  from typing import Optional
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from clu import platform
  from flax import jax_utils
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import ml_collections
  import optax
  from scenic.dataset_lib import dataset_utils
  from scenic.projects.baselines.bert import bert_base_model
  from scenic.projects.lang4video import util
  from scenic.projects.lang4video.model.image_text_model import ImageTextModel
  from scenic.projects.lang4video.trainer import optimizers
  from scenic.projects.lang4video.trainer.train_utils import ArgSpec
  from scenic.projects.lang4video.trainer.train_utils import axis_name_exists
  from scenic.projects.lang4video.trainer.train_utils import clip_grads
  from scenic.projects.lang4video.trainer.train_utils import compute_mask
  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 is_video_input
  from scenic.projects.lang4video.trainer.train_utils import NUM_DEVICES_AXIS_NAME
  from scenic.train_lib import lr_schedules
  from scenic.train_lib import pretrain_utils
  from scenic.train_lib import train_utils
label_desc: Use image node_id 2 for calling a global function with example usage: eval_step(train_state, visual, text, mask, batch_mask, text_for_mlm, segment_ids_for_mlm, mask_for_mlm, masked_lm_positions, masked_lm_ids, masked_lm_weights, text_batch_mask) and returns: metrics