code: def train(
      *,
      rng: jnp.ndarray,
      config: ml_collections.ConfigDict,
      model_cls: type[ImageTextModel],
      dataset: dataset_utils.Dataset,
      workdir: Optional[str] = None,  # pylint: disable=unused-argument
      writer: metric_writers.MetricWriter,
  ) -> tuple[Mapping[str, float], Mapping[str, float]]:
      """Trains a model for visual-text matching."""
      logging.info(
          "Creating the text dataset %s…", config.text_dataset_name
      )
      text_dataset = _create_text_dataset(config, rng, workdir)
      logging.info("Text dataset created.")
  
      input_spec = get_input_spec_for_pretraining(
          dataset_meta_data=dataset.meta_data,
          dataset_configs=config.get("dataset_configs", {}),
          train_=True,
          text_dataset_meta_data=text_dataset.meta_data,
      )
  
      is_video = is_video_input(input_spec[:3])
  
      lead_host = jax.process_index() == 0
  
      model = model_cls(config, dataset.meta_data)
  
      init_rng, rng = jax.random.split(rng)
      params, model_state = init_encoder(
          encoder=model.flax_model,
          input_spec=input_spec,
          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
          ),
          config=config,
          rng=init_rng,
      )
  
      lr_fn = lr_schedules.get_learning_rate_fn(config)
      optimizer_config = optimizers.get_optax_optimizer_config(config)
      tx = optimizers.get_optimizer(
          optimizer_config, lr_fn, params=params
      )
      if optimizer_config.get("lookahead"):
          params = optax.LookaheadParams.init_synced(params)
      opt_state = jax.jit(tx.init, backend="cpu")(params)
  
      chrono = train_utils.Chrono()
  
      train_rng, rng = jax.random.split(rng)
      train_state = train_utils.TrainState(
          tx=tx,
          opt_state=opt_state,
          params=params,
          model_state=model_state,
          rng=train_rng,
          metadata={"chrono": chrono.save()},
      )
  
      start_step = train_state.global_step
  
      if config.checkpoint:
          train_state, start_step = train_utils.restore_checkpoint(
              workdir, train_state
          )
  
      chrono.load(train_state.metadata["chrono"])
  
      if start_step == 0 and (init_from := config.get("init_from")):
          init_checkpoint_path = init_from.get("checkpoint_path")
          if init_checkpoint_path:
              restored_train_state = (
                  pretrain_utils.restore_pretrained_checkpoint(
                      init_checkpoint_path,
                      train_state,
                      assert_exist=True,
                  )
              )
              train_state = train_state.replace(
                  params=restored_train_state.params,
                  model_state=restored_train_state.model_state,
              )
              del restored_train_state
  
      train_state = jax_utils.replicate(train_state)
  
      if (batch_size := config.batch_size) == 1:
          logging.warning(
              "The batch size is 1, so then the NCE loss will be 0."
          )
  
      gather_scores = config.get("model", {}).get("gather_scores", True)
  
      train_step_pmapped = jax.pmap(
          functools.partial(
              train_step,
              model=model,
              is_video=is_video,
              lr_fn=lr_fn,
              config=config,
              gather_scores=gather_scores,
              debug=config.get("debug_train"),
          ),
          axis_name=NUM_DEVICES_AXIS_NAME,
          donate_argnums=tuple(range(10)),
      )
  
      eval_step_pmapped = jax.pmap(
          functools.partial(
              eval_step,
              model=model,
              is_video=is_video,
              gather_scores=gather_scores,
              debug=config.get("debug_eval"),
          ),
          axis_name=NUM_DEVICES_AXIS_NAME,
          donate_argnums=tuple(range(1, 12)),
      )
  
      total_steps, steps_per_epoch = train_utils.get_num_training_steps(
          config, dataset.meta_data
      )
      eval_steps = get_epoch_steps(config, dataset, split="eval")
  
      chrono.inform(
          first_step=start_step,
          total_steps=total_steps,
          global_bs=batch_size,
          steps_per_epoch=steps_per_epoch,
      )
      logging.info("Starting training loop at step %d.", start_step + 1)
  
      if jax.process_index() == 0:
          async_manager = (
              checkpoints.AsyncManager()
          )  # pylint: disable=unused-variable
      else:
          async_manager = None  # pylint: disable=unused-variable
  
      # We use this object to report the progress in the work unit notes (`chrono`
      # also writes similar ones, but whatever), but also to time the eval and
      # checkpoint saving. It writes some metrics to the writer, while `chrono`
      # writes others.
      report_progress = periodic_actions.ReportProgress(
          num_train_steps=total_steps, writer=writer
      )
  
      def write_note(note: str) -> None:
          if lead_host:
              platform.work_unit().set_notes(note)
  
      hooks = []
      if lead_host:
          hooks.append(report_progress)
      if lead_host and config.get("xprof", True):
          hooks.append(periodic_actions.Profile(logdir=workdir))
  
      train_summary, eval_summary = {}, {}
      metrics_list, training_logs_list = [], []
  
      write_note("Starting first step (and compilations)…")
      for step, batch, text_batch in zip(
          range(start_step + 1, total_steps + 1),
          dataset.train_iter,
          text_dataset.train_iter,
      ):
          with jax.profiler.StepTraceAnnotation("train", step_num=step):
              if (text := batch.get("text_indices")) is None:
                  text = batch["label"]
              else:
                  text = text[:, :, 0]
  
              mask = compute_mask(text, config)
  
              train_state, metrics, training_logs = train_step_pmapped(
                  train_state,
                  batch["inputs"],
                  text,
                  mask,
                  text_batch["input_word_ids"],
                  text_batch["input_type_ids"],
                  text_batch["input_mask"],
                  text_batch["masked_lm_positions"],
                  text_batch["masked_lm_ids"],
                  text_batch["masked_lm_weights"],
              )
  
              # This will accumulate metrics and training logs in the device memory up
              # to the point that we log them. This is no problem for small metrics but
              # may be a problem for large metrics. An alternative is to set
              # `log_summary_steps` to a small number, or to use
              # `train_utils.unreplicate_and_get` here instead of right before writing
              # summaries, but that means in each step, we have data transfer between
              # tpu and host, which might slow down the training.
              metrics_list.append(metrics)
              training_logs_list.append(training_logs)
  
          for h in hooks:
              h(step)
  
          if (
              (step % config.get("log_summary_steps", 100) == 0)
              or (step == 1)
              or (step == total_steps)
              or (lead_host and chrono.warmup)
          ):
              chrono.pause(wait_for=(metrics_list))
  
              if lead_host:
                  chrono.tick(step, writer, write_note)
  
              train_summary = train_utils.log_train_summary(
                  step=step,
                  writer=writer,
                  train_metrics=train_utils.unreplicate_and_get(
                      metrics_list
                  ),
                  extra_training_logs=train_utils.unreplicate_and_get(
                      training_logs_list
                  ),
                  key_separator="/",
                  flush_writer=False,
              )
              metrics_list, training_logs_list = [], []
  
              chrono.resume()
  
          if config.get("eval_while_training", True) and (
              (
                  step % config.get("log_eval_steps", steps_per_epoch)
                  == 0
              )
              or (step == 1)
              or (step == total_steps)
          ):
              chrono.pause(wait_for=train_state.params)
  
              with report_progress.timed("eval"):
                  eval_metrics_all = []
                  train_state = (
                      train_utils.sync_model_state_across_replicas(
                          train_state
                      )
                  )
                  for _, eval_batch, eval_text_batch in zip(
                      range(eval_steps),
                      dataset.valid_iter,
                      text_dataset.valid_iter,
                  ):
                      if (
                          text := eval_batch.get("text_indices")
                      ) is None:
                          text = eval_batch["label"]
                      else:
                          text = text[:, :, 0]
  
                      mask = compute_mask(text, config)
  
                      eval_metrics_batch = eval_step_pmapped(
                          train_state,
                          eval_batch["inputs"],
                          text,
                          mask,
                          eval_batch["batch_mask"],
                          eval_text_batch["input_word_ids"],
                          eval_text_batch["input_type_ids"],
                          eval_text_batch["input_mask"],
                          eval_text_batch["masked_lm_positions"],
                          eval_text_batch["masked_lm_ids"],
                          eval_text_batch["masked_lm_weights"],
                          eval_text_batch["batch_mask"],
                      )
                      eval_metrics_all.append(eval_metrics_batch)
                  eval_summary = train_utils.log_eval_summary(
                      step=step,
                      writer=writer,
                      eval_metrics=train_utils.unreplicate_and_get(
                          eval_metrics_all
                      ),
                      key_separator="/",
                      flush_writer=False,
                  )
  
              chrono.resume()
  
          if (
              (
                  step % config.get("checkpoint_steps", steps_per_epoch)
                  == 0
                  and step > 0
              )
              or (step == total_steps)
          ) and config.get("checkpoint"):
              chrono.pause(
                  wait_for=(train_state.params, train_state.opt_state)
              )
  
              with report_progress.timed("checkpoint"):
                  train_state = (
                      train_utils.sync_model_state_across_replicas(
                          train_state
                      )
                  )
                  if lead_host:
                      unreplicated_train_state = jax_utils.unreplicate(
                          train_state
                      )
                      metadata = unreplicated_train_state.metadata
                      metadata["chrono"] = chrono.save()
                      unreplicated_train_state = (
                          unreplicated_train_state.replace(
                              metadata=metadata
                          )
                      )  # pytype: disable=attribute-error
                      train_utils.save_checkpoint(
                          workdir,
                          unreplicated_train_state,
                          overwrite=config.get(
                              "overwrite_checkpoint", False
                          ),
                          # async_manager=async_manager,
                      )
                      del unreplicated_train_state
  
              chrono.resume()
  
      writer.flush()
  
      # Wait until computations are done before exiting.
      train_utils.barrier_across_hosts()
  
      return train_summary, eval_summary
  
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 5 for calling a global function with example usage: train() and returns: train_summary,  eval_summary