code: def train_and_eval(
      rng: np.ndarray,
      config: ml_collections.ConfigDict,
      *,
      workdir: str,
      writer: Any,
      model_cls,
      dataset,
  ) -> Tuple[utils.OptaxTrainState, 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: 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
  
      model = model_cls(config, dataset.meta_data)
      train_step_pmapped, eval_step_pmapped = pmapped_steps(
          model, config
      )
      train_state, start_step, chrono = utils.init_state(
          model,
          dataset,
          config,  # pytype: disable=wrong-arg-types  # jax-ndarray
          workdir,
          rng,
      )
      chrono.load(train_state.metadata["chrono"])
      train_state = jax_utils.replicate(train_state)
  
      # Calculate the total number of training steps.
      total_steps, steps_per_epoch = train_utils.get_num_training_steps(
          config, dataset.meta_data
      )
      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
      )
  
      # And the number of evaluation steps.
      num_eval_examples = dataset.meta_data["num_eval_examples"]
      logging.info("Number of processes is %s", jax.process_count())
      total_eval_steps = int(
          np.ceil(num_eval_examples / (config.get("batch_size")))
      )
      steps_per_eval = config.get("steps_per_eval", total_eval_steps)
      train_metrics = []
      train_summary = None
  
      eval_summary = {}
      logging.info("Start training from step %d", start_step + 1)
      report_progress = periodic_actions.ReportProgress(
          num_train_steps=total_steps, writer=writer
      )
  
      def write_note(note):
          if lead_host:
              platform.work_unit().set_notes(note)
  
      hooks = []
      if lead_host:
          hooks.append(report_progress)
      if config.get("xprof", True) and jax.process_index() == 0:
          hooks.append(
              periodic_actions.Profile(
                  num_profile_steps=5, logdir=workdir
              )
          )
      chrono.inform(
          start_step, total_steps, config.batch_size, steps_per_epoch
      )
      for step in range(start_step + 1, total_steps + 1):
          with jax.profiler.StepTraceAnnotation("train", step_num=step):
              train_batch = next(dataset.train_iter)
              train_state, t_metrics = train_step_pmapped(
                  train_state, train_batch
              )
              train_metrics.append(t_metrics)
          for hook in hooks:
              hook(step)
          # Log the train summary every `log_summary_steps`.
          if (step % log_summary_steps == 1) or (step == total_steps):
              chrono.pause()
              if lead_host:
                  chrono.tick(step, writer, write_note)
              train_summary = train_utils.log_train_summary(
                  step=step,
                  train_metrics=jax.tree_util.tree_map(
                      train_utils.unreplicate_and_get, train_metrics
                  ),
                  writer=writer,
              )
              train_metrics = []
              chrono.resume()
  
          # Evaluate every `log_eval_steps`.
          do_eval = (step % log_eval_steps == 1) or (
              step == total_steps
          )
          if do_eval:
              eval_summary = eval_and_log_summary(
                  train_state=train_state,
                  iterator=dataset.valid_iter,
                  eval_step_fn=eval_step_pmapped,
                  eval_steps=steps_per_eval,
                  writer=writer,
                  train_iteration=step,
                  num_eval_examples=num_eval_examples,
                  compute_recall_metrics=True,
              )
              writer.flush()
  
          # Checkpointing.
          if not config.checkpoint:
              continue
          elif do_eval or (step % checkpoint_steps == 0 and step > 0):
              chrono.pause(
                  wait_for=(train_state.weights, train_state.opt_state)
              )
              if lead_host:
                  # Take the first replica.
                  unrep_train_state = jax_utils.unreplicate(train_state)
                  metadata = unrep_train_state.metadata
                  metadata["chrono"] = chrono.save()
                  unrep_train_state.replace(
                      metadata=metadata
                  )  # pytype: disable=attribute-error
                  utils.save_checkpoint(workdir, unrep_train_state)
                  del unrep_train_state
              chrono.resume()  # Un-pause now.
      # Wait until computations are done before exiting.
      train_utils.barrier_across_hosts()
      # Return the train and eval summary after last step for regression testing.
      return train_state, train_summary, eval_summary
  
imports: import functools
  from typing import Any, Dict, Optional, Tuple
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from clu import platform
  from flax import jax_utils
  import flax.linen
  import jax
  import jax.example_libraries.optimizers
  import jax.profiler
  import ml_collections
  import numpy
  import optax
  from scenic.projects.verbs_in_action import utils
  from scenic.train_lib import train_utils
label_desc: Use image node_id 1 for calling a global function with example usage: train_and_eval(rng, config) and returns: train_state,  train_summary,  eval_summary