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],  # Shape: (N, L)
      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,
      *,
      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, jnp.ndarray]
      ]:
          (
              encoded_visual,
              encoded_text,
              mlm_logits,
          ), 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_and_do_pretraining
                  if is_video
                  else model.flax_model.encode_image_and_text_and_do_pretraining
              ),
              rngs={"dropout": rng},
              **{"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=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_image_text_ = model.loss_function(scores).mean()
          loss_mlm_ = (
              bert_base_model.sparse_weighted_softmax_cross_entropy(
                  mlm_logits, masked_lm_ids, masked_lm_weights
              )
          )
          loss_ = (
              loss_image_text_
              + config.get("mlm_loss_weight", 1) * loss_mlm_
          )
  
          return loss_, (
              new_model_state_,
              loss_image_text_,
              loss_mlm_,
          )  # pytype: disable=bad-return-type  # jax-ndarray
  
      grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
      (
          loss,
          (new_model_state, loss_image_text, loss_mlm),
      ), 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, 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,
          )
      )
  
      metrics = {
          "loss": loss,
          "loss/image_text": loss_image_text,
          "loss/mlm": loss_mlm,
      }
  
      # We don't need to look at `batch_mask` because we drop the remainder.
      local_batch_size = len(visual)
      for k, v in metrics.items():
          metrics[k] = (v * 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
  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 1 for calling a global function with example usage: train_step(train_state, visual, text, mask, text_for_mlm, segment_ids_for_mlm, mask_for_mlm, masked_lm_positions, masked_lm_ids, masked_lm_weights) and returns: new_train_state,  metrics,  training_logs, loss_,  new_model_state_,  loss_image_text_,  loss_mlm_