code: def train_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,
      lr_fn: Callable[[int], float],
      config: ml_collections.ConfigDict,
      gather_scores: bool = True,
      debug: bool = False,
  ) -> tuple[
      train_utils.TrainState, Mapping[str, Any], Mapping[str, Any]
  ]:
      """Runs a single step of evaluation."""
      encoder = model.flax_model
  
      new_rng, rng = jax.random.split(train_state.rng)
  
      if axis_name_exists(NUM_DEVICES_AXIS_NAME):
          rng = train_utils.bind_rng_to_host_device(
              rng, NUM_DEVICES_AXIS_NAME, bind_to="device"
          )
  
      def loss_fn(
          params: Mapping[str, Any]
      ) -> tuple[float, tuple[Mapping[str, Any], jnp.ndarray]]:
          (
              encoded_visual,
              encoded_text,
          ), new_model_state_ = encoder.apply(
              variables={
                  "params": params,
                  **train_state.model_state,
              },
              mutable=train_state.model_state.keys(),
              method=model.flax_model.encode_video_and_text
              if is_video
              else None,
              rngs={"dropout": rng},
              **{"video" if is_video else "image": visual},
              text=text,
              mask=mask,
              train=True,
              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
              ),
          )
  
          loss_ = model.loss_function(scores).mean()
  
          return (
              loss_,
              new_model_state_,
          )  # pytype: disable=bad-return-type  # jax-ndarray
  
      grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
      (loss, new_model_state), grad = grad_fn(train_state.params)
  
      if axis_name_exists(NUM_DEVICES_AXIS_NAME):
          grad = jax.lax.pmean(grad, axis_name=NUM_DEVICES_AXIS_NAME)
  
      if max_grad_norm := config.get("max_grad_norm"):
          grad = clip_grads(grad, max_grad_norm)
  
      updates, new_opt_state = train_state.tx.update(
          grad,
          train_state.opt_state,  # pytype: disable=attribute-error
          train_state.params,
      )
      new_params = optax.apply_updates(train_state.params, updates)
  
      new_train_state = (
          train_state.replace(  # pytype: disable=attribute-error
              global_step=train_state.global_step + 1,
              opt_state=new_opt_state,
              model_state=new_model_state,
              params=new_params,
              rng=new_rng,
          )
      )
  
      # We don't need to look at `batch_mask` because we drop the remainder.
      local_batch_size = len(visual)
      metrics = {"loss": (loss * local_batch_size, local_batch_size)}
  
      training_logs = {
          "temperature": 1 / jnp.exp(encoder.logit_scale),
          "step": train_state.global_step,
          "batch_size": local_batch_size * jax.device_count(),
          "lr": lr_fn(train_state.global_step),
      }
  
      return new_train_state, metrics, training_logs
  
imports: from collections.abc import Callable
  from collections.abc import Mapping
  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.lang4video.model.image_text_model import ImageTextModel
  from scenic.projects.lang4video.trainer import optimizers
  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 1 for calling a global function with example usage: train_step(train_state, visual, text, mask) and returns: new_train_state,  metrics,  training_logs, loss_,  new_model_state_