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] = None,  # Shape: (N, L)
      *,
      model: ImageTextModel,
      is_video: bool,
      debug: bool = False,
  ) -> tuple[jnp.ndarray, jnp.ndarray]:  # Shapes: (N, E), (N, E)
      """Runs a single step of evaluation."""
      encoded_visual, encoded_text = model.flax_model.apply(
          variables={
              "params": train_state.params,
              **train_state.model_state,
          },
          method=model.flax_model.encode_video_and_text
          if is_video
          else None,
          **{"video" if is_video else "image": visual},
          text=text,
          mask=mask,
          train=False,
          debug=debug,
      )
  
      if axis_name_exists(NUM_DEVICES_AXIS_NAME):
          encoded_visual = jax.lax.all_gather(
              encoded_visual, NUM_DEVICES_AXIS_NAME
          )
          encoded_visual = encoded_visual.reshape(
              -1, encoded_visual.shape[-1]
          )
  
          encoded_text = jax.lax.all_gather(
              encoded_text, NUM_DEVICES_AXIS_NAME
          )
          encoded_text = encoded_text.reshape(
              -1, encoded_text.shape[-1]
          )
  
      return encoded_visual, encoded_text
  
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 1 for calling a global function with example usage: eval_step(train_state, visual, text, mask) and returns: encoded_visual,  encoded_text