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.
      target: jnp.ndarray,  # Shape: (N) as int or (N, K) as float
      batch_mask: jnp.ndarray,  # Shape: (N,)
      *,
      model: ImageTextModel,
      encoded_classes: jnp.ndarray,  # Shape: (K, E)
      is_video: bool,
      debug: bool = False,
  ) -> Mapping[str, tuple[float, int]]:
      """Runs a single step of evaluation."""
      encoded_visual = model.flax_model.apply(
          variables={
              "params": train_state.params,
              **train_state.model_state,
          },
          method=model.flax_model.encode_video
          if is_video
          else model.flax_model.encode_image,
          **{"video" if is_video else "image": visual},
          train=False,
          debug=debug,
      )
  
      score = model.flax_model.compute_similarity(
          encoded_visual, encoded_classes
      )
  
      if target.ndim == score.ndim - 1:
          target = jax.nn.one_hot(target, len(encoded_classes))
      assert target.ndim == score.ndim
  
      metrics = {
          "accuracy@1": model_utils.weighted_top_one_correctly_classified(
              score, target, weights=batch_mask
          ),
          "accuracy@5": model_utils.weighted_topk_correctly_classified(
              score, target, k=5, weights=batch_mask
          ),
          "loss": optax.softmax_cross_entropy(score, target)
          * batch_mask,
      }
  
      actual_local_batch_size = batch_mask.sum()
      for k, v in metrics.items():
          metrics[k] = (v, actual_local_batch_size)
  
      if axis_name_exists(NUM_DEVICES_AXIS_NAME):
          actual_batch_size = jax.lax.psum(
              actual_local_batch_size, 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,
              )
  
      return metrics
  
imports: from collections.abc import Mapping
  import dataclasses
  import functools
  from typing import Optional
  from absl import logging
  from clu import metric_writers
  from flax import jax_utils
  import flax.core
  import jax
  import jax.numpy
  import ml_collections
  import optax
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.lang4video import util
  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_array_to_be_divisible
  from scenic.projects.lang4video.trainer.train_utils import partial_with_cache
  from scenic.projects.lang4video.trainer.train_utils import split_in_batches
  from scenic.train_lib import train_utils
  from tqdm.auto import tqdm
label_desc: Use image node_id 1 for calling a global function with example usage: eval_step(train_state, visual, target, batch_mask) and returns: metrics