code: def train_and_evaluate(
      *,
      rng: jnp.ndarray,
      config: ml_collections.ConfigDict,
      model_cls: Any,
      dataset: dataset_utils.Dataset,
      workdir: str,
      writer: metric_writers.MetricWriter,
  ) -> Tuple[train_utils.TrainState, Dict[str, Any], Dict[str, Any]]:
      """Main training loop lives in this function.
  
      Given the model class and dataset, it prepares the items needed to run the
      training, including the TrainState.
  
      Args:
        rng: JAX PRNGKey.
        config: Configurations of the experiment.
        model_cls: Model class; A model has a flax_module, a loss_fn, and a
          metrics_fn associated with it.
        dataset: The dataset that has train_iter, eval_iter, meta_data, and
          optionally, test_iter.
        workdir: Directory for checkpointing.
        writer: CLU metrics writer instance.
  
      Returns:
        train_sate that has the state of training (including current
          global_step, model_state, rng, and the optimizer), train_summary
          and eval_summary which are dict of metrics. These outputs are used for
          regression testing.
      """
      lead_host = jax.process_index() == 0
  
      # The pool is used to perform misc operations such as logging in async way.
      pool = futures.ThreadPoolExecutor(max_workers=2)
  
      # Build the loss_and_metrics_fn, metrics, and flax_model.
      model = model_cls(config, dataset.meta_data)
  
      # Initialize model.
      rng, init_rng = jax.random.split(rng)
      (
          params,
          model_state,
          num_trainable_params,
          gflops,
      ) = train_utils.initialize_model(
          model_def=model.flax_model,
          input_spec=[
              (
                  dataset.meta_data["input_shape"],
                  dataset.meta_data.get("input_dtype", jnp.float32),
              )
          ],
          config=config,
          rngs=init_rng,
      )
  
      # Create optimizer.
      # We jit this, such that the arrays that are created are created on the same
      # device as the input is, in this case the CPU. Else they'd be on device[0]
      optimizer = jax.jit(
          detr_train_utils.get_detr_optimizer(config).create,
          backend="cpu",
      )(params)
  
      rng, train_rng = jax.random.split(rng)
      train_state = train_utils.TrainState(
          global_step=0,
          optimizer=optimizer,
          model_state=model_state,
          rng=train_rng,
          accum_train_time=0,
      )
      start_step = train_state.global_step
      if config.checkpoint:
          train_state, start_step = train_utils.restore_checkpoint(
              workdir, train_state
          )
  
      if (
          start_step == 0  # Which means "no" checkpoint is restored!
          and config.get("init_from") is not None
      ):
          init_checkpoint_path = config.init_from.get("checkpoint_path")
          restored_train_state = flax_restore_checkpoint(
              init_checkpoint_path, target=None
          )
          train_state = pretrain_utils.init_from_pretrain_state(
              train_state,
              restored_train_state,
              ckpt_prefix_path=config.init_from.get("ckpt_prefix_path"),
              model_prefix_path=config.init_from.get(
                  "model_prefix_path"
              ),
              name_mapping=config.init_from.get("name_mapping"),
              skip_regex=config.init_from.get("skip_regex"),
          )
          # Free unecessary memory.
          del restored_train_state
      elif start_step == 0 and config.get(
          "load_pretrained_backbone", False
      ):
          # Only load pretrained backbone if we are at the beginning of training.
          bb_checkpoint_path = config.pretrained_backbone_configs.get(
              "checkpoint_path"
          )
          bb_train_state = flax_restore_checkpoint(
              bb_checkpoint_path, target=None
          )
  
          model_prefix_path = ["backbone"]
          train_state = pretrain_utils.init_from_pretrain_state(
              train_state,
              bb_train_state,
              model_prefix_path=model_prefix_path,
          )
  
      update_model_state = not config.get(
          "freeze_backbone_batch_stats", False
      )
      if not update_model_state:
          if not config.load_pretrained_backbone:
              raise ValueError(
                  "Freezing the batch statistics of the resnet backbone "
                  "is only possible when loading a pretrained resnet "
                  "backbone is enabled."
              )
      # Replicate the optimzier, state, and rng.
      train_state = jax_utils.replicate(train_state)
      del params  # Do not keep a copy of the initial params.
  
      # Calculate the total number of training steps.
      total_steps, steps_per_epoch = train_utils.get_num_training_steps(
          config, dataset.meta_data
      )
      # Get learning rate scheduler.
      learning_rate_fn = lr_schedules.get_learning_rate_fn(config)
      backbone_learning_rate_fn = None
      if config.get("backbone_training"):
          backbone_learning_rate_fn = lr_schedules.get_learning_rate_fn(
              config.backbone_training
          )
  
      train_step = get_train_step(
          flax_model=model.flax_model,
          loss_and_metrics_fn=model.loss_function,
          learning_rate_fn=learning_rate_fn,
          backbone_learning_rate_fn=backbone_learning_rate_fn,
          max_grad_norm=config.get("max_grad_norm", None),
          update_model_state=update_model_state,
          debug=config.debug_train,
      )
  
      train_step_pmapped = jax.pmap(
          train_step, axis_name="batch", donate_argnums=(0,)
      )
  
      ############### EVALUATION CODE #################
      eval_step = get_eval_step(
          flax_model=model.flax_model,
          loss_and_metrics_fn=model.loss_function,
          logits_to_probs_fn=model.logits_to_probs,
          debug=config.debug_eval,
      )
      eval_step_pmapped = jax.pmap(
          eval_step, axis_name="batch", donate_argnums=(1,)
      )
  
      # Ceil rounding such that we include the last incomplete batch.
      eval_batch_size = config.get("eval_batch_size", config.batch_size)
      total_eval_steps = int(
          np.ceil(
              dataset.meta_data["num_eval_examples"] / eval_batch_size
          )
      )
      steps_per_eval = config.get("steps_per_eval") or total_eval_steps
  
      metrics_normalizer_fn = functools.partial(
          detr_train_utils.normalize_metrics_summary,
          object_detection_loss_keys=model.loss_terms_weights.keys(),
      )
  
      def evaluate(train_state, step):
          """Runs evaluation code."""
          future = None
  
          def _wait(
              future: Optional[futures.Future],
          ) -> Any:  # pylint: disable=g-bare-generic
              if future is None:
                  return None
              return future.result()
  
          def _add_examples(predictions, labels):
              for pred, label in zip(predictions, labels):
                  global_metrics_evaluator.add_example(
                      prediction=pred, target=label
                  )
  
          eval_metrics = []
          if global_metrics_evaluator is not None:
              global_metrics_evaluator.clear()
  
          for eval_step in range(steps_per_eval):
              logging.info("Running eval step %d", eval_step)
              eval_batch = next(dataset.valid_iter)
  
              # Do the eval step given the matches.
              (
                  eval_batch_all_hosts,
                  eval_predictions_all_hosts,
                  e_metrics,
              ) = eval_step_pmapped(train_state, eval_batch)
  
              # Variable aux_outputs is not needed anymore.
              eval_predictions_all_hosts.pop("aux_outputs", None)
  
              # Collect local metrics (returned by the loss function).
              eval_metrics.append(
                  train_utils.unreplicate_and_get(e_metrics)
              )
  
              if global_metrics_evaluator is not None:
                  # Unreplicate the output of eval_step_pmapped (used `lax.all_gather`).
                  eval_batch_all_hosts = jax_utils.unreplicate(
                      eval_batch_all_hosts
                  )
                  eval_predictions_all_hosts = jax_utils.unreplicate(
                      eval_predictions_all_hosts
                  )
  
                  # Collect preds and labels to be sent for computing global metrics.
                  predictions = (
                      detr_train_utils.process_and_fetch_to_host(
                          eval_predictions_all_hosts,
                          eval_batch_all_hosts["batch_mask"],
                      )
                  )
                  predictions = jax.tree_util.tree_map(
                      np.asarray, predictions
                  )
  
                  labels = detr_train_utils.process_and_fetch_to_host(
                      eval_batch_all_hosts["label"],
                      eval_batch_all_hosts["batch_mask"],
                  )
                  labels = jax.tree_util.tree_map(np.asarray, labels)
  
                  if eval_step == 0:
                      logging.info(
                          "Pred keys: %s", list(predictions[0].keys())
                      )
                      logging.info(
                          "Labels keys: %s", list(labels[0].keys())
                      )
  
                  # Add to evaluator.
                  _wait(future)
                  future = pool.submit(
                      _add_examples, predictions, labels
                  )
  
                  del predictions, labels
  
              del (
                  eval_batch,
                  eval_batch_all_hosts,
                  eval_predictions_all_hosts,
              )
  
          eval_global_metrics_summary_future = None
          if global_metrics_evaluator is not None:
              _wait(future)
              logging.info(
                  "Number of eval examples: %d",
                  len(global_metrics_evaluator),
              )
              if lead_host:
                  eval_global_metrics_summary_future = pool.submit(
                      global_metrics_evaluator.compute_metrics,
                      clear_annotations=False,
                  )
  
          return (
              step,
              eval_metrics,
          ), eval_global_metrics_summary_future
  
      ###################################################
  
      log_eval_steps = config.get("log_eval_steps") or steps_per_epoch
      if not log_eval_steps:
          raise ValueError(
              "'log_eval_steps' should be specified in the config."
          )
      log_summary_steps = config.get("log_summary_steps", 25)
      log_large_summary_steps = config.get("log_large_summary_steps", 0)
      checkpoint_steps = (
          config.get("checkpoint_steps") or log_eval_steps
      )
  
      global_metrics_evaluator = (
          None  # Only run eval on the lead_host node.
      )
      if lead_host:
          global_metrics_evaluator = (
              detr_train_utils.DetrGlobalEvaluator(config.dataset_name)
          )
  
      train_metrics, extra_training_logs = [], []
      train_summary, eval_summary = None, None
  
      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("Starting training loop at step %d.", start_step + 1)
      report_progress = periodic_actions.ReportProgress(
          num_train_steps=total_steps,
          writer=writer,
          every_secs=None,
          every_steps=config.get(
              "report_progress_step", log_summary_steps
          ),
      )
      hooks = []
      if lead_host:
          hooks.append(report_progress)
      if config.get("xprof", True) and lead_host:
          hooks.append(
              periodic_actions.Profile(
                  num_profile_steps=5, logdir=workdir
              )
          )
  
      if start_step == 0:
          step0_log = {"num_trainable_params": num_trainable_params}
          if gflops:
              step0_log["gflops"] = gflops
          writer.write_scalars(1, step0_log)
  
      (last_eval_step, last_eval_metrics), last_eval_future = (
          None,
          None,
      ), None
      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,
                  lr,
                  backbone_lr,
                  train_predictions,
              ) = train_step_pmapped(train_state, train_batch)
              # This will accumulate metrics in TPU memory up to the point that we log
              # them. This is no problem for small metrics but may be a problem for
              # large (e.g. segmentation) 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.
              train_metrics.append(t_metrics)
              # Additional training logs: learning rate, learning_rate_backbone:
              extra_training_logs.append(
                  {
                      "learning_rate": lr,
                      "learning_rate_backbone": backbone_lr,
                  }
              )
  
          for h in hooks:
              h(step)
  
          chrono.pause()
          if (
              log_large_summary_steps
              and step % log_large_summary_steps == 0
              and lead_host
          ):
              ############### LOG EXPENSIVE TRAIN SUMMARY ###############
              # Visualizes detections using side-by-side gt-pred images.
              # TODO(mjlm): Investigate this error when including `batch_mask`:
              # RuntimeError: Invalid argument: from_python argument must be an array.
              to_cpu = lambda x: jax.device_get(
                  dataset_utils.unshard(x)
              )
              del train_batch["batch_mask"]
              train_pred_cpu = to_cpu(train_predictions)
              train_batch_cpu = to_cpu(train_batch)
              viz = detr_train_utils.draw_boxes_side_by_side(
                  train_pred_cpu,
                  train_batch_cpu,
                  label_map=dataset.meta_data["label_to_name"],
              )
              viz_detections = {
                  f"sidebyside_{i}/detection": viz_[None, ...]
                  for i, viz_ in enumerate(viz)
              }
              writer.write_images(step, viz_detections)
  
          del train_predictions
  
          if (step % log_summary_steps == 0) or (
              step == total_steps - 1
          ):
              ############### LOG TRAIN SUMMARY ###############
              if lead_host:
                  chrono.tick(step, writer)
  
              # Write summary:
              train_summary = train_utils.log_train_summary(
                  step=step,
                  train_metrics=jax.tree_util.tree_map(
                      train_utils.unreplicate_and_get, train_metrics
                  ),
                  extra_training_logs=jax.tree_util.tree_map(
                      train_utils.unreplicate_and_get,
                      extra_training_logs,
                  ),
                  writer=writer,
                  metrics_normalizer_fn=metrics_normalizer_fn,
              )
              # Reset metric accumulation for next evaluation cycle.
              train_metrics, extra_training_logs = [], []
              #################################################
  
          if (step % log_eval_steps == 0) or (step == total_steps):
              # First wait for the previous eval to finish & write summary.
              if last_eval_future is not None:
                  train_utils.log_eval_summary(
                      step=last_eval_step,
                      eval_metrics=last_eval_metrics,
                      extra_eval_summary=last_eval_future.result(),
                      writer=writer,
                      metrics_normalizer_fn=metrics_normalizer_fn,
                  )
                  last_eval_future = None
  
              # Sync model state across replicas (in case of having model state, e.g.
              # batch statistic when using batch norm).
              start_time = time.time()
              with report_progress.timed("eval"):
                  train_state = (
                      train_utils.sync_model_state_across_replicas(
                          train_state
                      )
                  )
                  (
                      last_eval_step,
                      last_eval_metrics,
                  ), last_eval_future = evaluate(train_state, step)
              duration = time.time() - start_time
              logging.info(
                  "Done with async evaluation: %.4f sec.", duration
              )
              writer.flush()
  
          ##################### CHECKPOINTING ############################
          if (
              (step % checkpoint_steps == 0 and step > 0)
              or (step == total_steps)
          ) and config.checkpoint:
              with report_progress.timed("checkpoint"):
                  # Sync model state across replicas.
                  train_state = (
                      train_utils.sync_model_state_across_replicas(
                          train_state
                      )
                  )
                  if lead_host:
                      train_state.replace(
                          accum_train_time=chrono.accum_train_time
                      )
                      train_utils.save_checkpoint(workdir, train_state)
  
          chrono.resume()  # Un-pause now.
  
      # Wait until computations are done before exiting.
      pool.shutdown()
      if last_eval_future is not None:
          train_utils.log_eval_summary(
              step=last_eval_step,
              eval_metrics=last_eval_metrics,
              extra_eval_summary=last_eval_future.result(),
              writer=writer,
              metrics_normalizer_fn=metrics_normalizer_fn,
          )
      jax.random.normal(jax.random.PRNGKey(0), ()).block_until_ready()
      return train_state, train_summary, eval_summary
  
imports: from concurrent import futures
  import functools
  import time
  from typing import Any, Dict, Tuple, Optional
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from flax import jax_utils
  from flax.training.checkpoints import restore_checkpoint
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.projects.baselines.detr import train_utils
  from scenic.train_lib_deprecated import lr_schedules
  from scenic.train_lib_deprecated import pretrain_utils
  from scenic.train_lib_deprecated import train_utils
label_desc: Use image node_id 3 for calling a global function with example usage: train_and_evaluate() and returns: train_state,  train_summary,  eval_summary, step,  eval_metrics,  eval_global_metrics_summary_future, future, None