code: def train_and_eval(
      rng: np.ndarray,
      config: ml_collections.ConfigDict,
      *,
      workdir: str,
      writer: Any,
      model_cls,
      dataset_dict,
  ) -> Tuple[train_utils.TrainState, Dict[str, Any], Dict[str, Any]]:
      """Train (and occasionally evaluate) the model.
  
      Args:
        rng: JAX prng key.
        config: The configuration of the experiment.
        workdir: Where to checkpoint and write the summaries.
        writer: Summary writer object.
        model_cls: The model class used to instantiate the model.
        dataset_dict: The dataset for training and evaluation.
  
      Returns:
        A tuple with:
          * the state that has the state of training (including current
            global_step, model_state, rng, and the optimizer)
          * a dictionary with the train_summary
          * a dictionary with the evaluation summary
      """
  
      lead_host = jax.host_id() == 0
  
      datasets_metadata = {
          name: ds.meta_data for name, ds in dataset_dict.items()
      }
      all_datasets = []
      all_datasets_num_train_examples = []
      for name, metadata in datasets_metadata.items():
          all_datasets.append(name)
          all_datasets_num_train_examples.append(
              metadata.get("num_train_examples", 0)
          )
      model = model_cls(config, datasets_metadata)
      train_step_pmapped, eval_step_pmapped = pmapped_steps(
          model, config
      )
  
      train_state, start_step = init_state(
          model,
          dataset_dict[
              all_datasets[0]
          ],  # pytype: disable=wrong-arg-types  # jax-ndarray
          config,
          workdir,
          rng,
      )
      train_state = jax_utils.replicate(train_state)
      logging.info("Number of processes is %s", jax.process_count())
  
      del rng  # So that we don't mistakenly re-use it.
  
      # Calculate the total number of training steps.
      (
          total_steps,
          steps_per_epoch,
      ) = vid2seq_train_utils.get_num_training_steps_multi(
          config, datasets_metadata
      )
      log_eval_steps = config.get("log_eval_steps", steps_per_epoch)
      checkpoint_steps = config.get("checkpoint_steps", log_eval_steps)
      log_summary_steps = config.get(
          "log_summary_steps", log_eval_steps
      )
  
      # Build a tokenizer for the evaluation
      tokenizer = get_tokenizer(config)
  
      chrono = train_utils.Chrono(
          first_step=start_step,
          total_steps=total_steps,
          steps_per_epoch=steps_per_epoch,
          global_bs=config.batch_size,
          accum_train_time=int(
              jax_utils.unreplicate(train_state.accum_train_time)
          ),
      )
  
      logging.info("Start training from step %d", start_step + 1)
      hooks = []
      if config.get("xprof", True) and jax.process_index() == 0:
          hooks.append(
              periodic_actions.Profile(
                  num_profile_steps=5, logdir=workdir
              )
          )
  
      summary_builder = {
          x: SummaryBuilder([], []) for x in all_datasets
      }
      train_summary, eval_summary = None, None
  
      def get_next_train_batch(all_datasets, step):
          dataset = random.Random(step).choices(
              all_datasets,
              config.get(
                  "probs", [1.0 / len(all_datasets)] * len(all_datasets)
              ),
          )[0]
          ds = dataset_dict[dataset]
  
          return next(ds.train_iter), dataset
  
      for step in range(start_step + 1, total_steps + 1):
          chrono.resume()
          train_batch, train_ds = get_next_train_batch(
              all_datasets, step
          )
          with jax.profiler.StepTraceAnnotation("train", step_num=step):
              (
                  train_state,
                  train_loss,
                  t_metrics,
                  lr,
              ) = train_step_pmapped(train_ds, train_state, train_batch)
          for hook in hooks:
              hook(step)
          summary_builder[train_ds].update(
              t_metrics, {"lr": lr, "train/loss": train_loss}
          )
          chrono.pause()
  
          # Log the train summary every `log_summary_steps`.
          if (step % log_summary_steps == 1) or (step == total_steps):
              if lead_host:
                  chrono.tick(step, writer)
              train_summary = {x: {} for x in all_datasets}
              for x in all_datasets:
                  if len(
                      summary_builder[x].metrics
                  ):  # pylint: disable=g-explicit-length-test
                      train_summary[x].update(
                          summary_builder[x].write(writer, step)
                      )
  
          # Evaluate every `log_eval_steps`.
          should_eval = (step % log_eval_steps == 1) or (
              step == total_steps
          )
          if should_eval:
              for ds_name in dataset_dict[all_datasets[0]].valid_iter:
                  # Compute the number of evaluation steps per dataset.
                  num_eval_examples = dataset_dict[
                      all_datasets[0]
                  ].meta_data["num_eval_examples"][ds_name]
                  total_eval_steps = int(
                      np.ceil(
                          num_eval_examples
                          / (config.get("eval_batch_size"))
                      )
                  )
                  steps_per_eval = config.get(
                      "steps_per_eval", total_eval_steps
                  )
  
                  eval_summary = eval_and_log_summary(
                      train_state=train_state,
                      iterator=dataset_dict[all_datasets[0]].valid_iter[
                          ds_name
                      ],
                      eval_step_fn=eval_step_pmapped,
                      eval_steps=steps_per_eval,
                      writer=writer,
                      train_iteration=step,
                      tokenizer=tokenizer,
                      dataset_name=ds_name,
                      num_bins=config.dataset_configs.num_bins,
                      vocabulary_size=config.dataset_configs.vocabulary_size,
                      abs_time_token=config.dataset_configs.abs_time_token,
                      time_format=config.dataset_configs.time_format,
                      tmp_only=config.dataset_configs.tmp_only,
                      runlocal=config.runlocal,
                      order=config.dataset_configs.order,
                      workdir=workdir,
                      soda="soda" in config.eval_metrics,
                      eval_batch_size=config.eval_batch_size,
                      max_events=config.dataset_configs.max_events,
                      para="para" in ds_name,
                      t=1000000.0,  # 1 FPS
                      is_split=config.dataset_configs.split,
                  )
  
          # Checkpoint.
          if not config.checkpoint:
              continue
          elif step % checkpoint_steps == 0 and step > 0:
              logging.info(
                  "checkpointing starts (training step: %d)", step
              )
              # Sync model state across replicas.
              train_state = (
                  train_utils.sync_model_state_across_replicas(
                      train_state
                  )
              )
              logging.info("checkpointing (training step: %d)", step)
              train_state.replace(  # pytype: disable=attribute-error
                  accum_train_time=chrono.accum_train_time
              )
              train_utils.save_checkpoint(
                  workdir, train_state, max_to_keep=100
              )
  
      # Return the train and eval summary after last step.
      return train_state, train_summary, eval_summary
  
imports: import copy
  import dataclasses
  import functools
  import json
  import os
  import random
  from typing import Any, Callable, Dict, List, Optional, Tuple
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from dmvr import tokenizers
  from flax import jax_utils
  from flax.core import unfreeze
  import flax.linen
  import jax
  import jax.example_libraries.optimizers
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import base_model
  from scenic.projects.t5 import model
  from scenic.projects.t5 import tokenizer
  from scenic.projects.vid2seq import load_utils
  from scenic.projects.vid2seq import models
  from scenic.projects.vid2seq import train_utils
  from scenic.projects.vid2seq import dvc_eval
  from scenic.train_lib_deprecated import lr_schedules
  from scenic.train_lib_deprecated import optimizers
  from scenic.train_lib_deprecated import pretrain_utils
  from scenic.train_lib_deprecated import train_utils
label_desc: Use image node_id 12 for calling a global function with example usage: train_and_eval(rng, config) and returns: train_state,  train_summary,  eval_summary, next,  dataset