code: def train_and_evaluate(
      config: ml_collections.ConfigDict, workdir: str
  ) -> TrainState:
      """Execute model training and evaluation loop.
  
      Args:
        config: Hyperparameter configuration for training and evaluation.
        workdir: Directory where the tensorboard summaries are written to.
  
      Returns:
        Final TrainState.
      """
  
      writer = metric_writers.create_default_writer(
          logdir=workdir, just_logging=jax.process_index() != 0
      )
  
      rng = random.key(0)
  
      image_size = 224
  
      if config.batch_size % jax.device_count() > 0:
          raise ValueError(
              "Batch size must be divisible by the number of devices"
          )
      local_batch_size = config.batch_size // jax.process_count()
  
      platform = jax.local_devices()[0].platform
  
      if config.half_precision:
          if platform == "tpu":
              input_dtype = tf.bfloat16
          else:
              input_dtype = tf.float16
      else:
          input_dtype = tf.float32
  
      dataset_builder = tfds.builder(config.dataset)
      train_iter = create_input_iter(
          dataset_builder,
          local_batch_size,
          image_size,
          input_dtype,
          train=True,
          cache=config.cache,
          shuffle_buffer_size=config.shuffle_buffer_size,
          prefetch=config.prefetch,
      )
      eval_iter = create_input_iter(
          dataset_builder,
          local_batch_size,
          image_size,
          input_dtype,
          train=False,
          cache=config.cache,
          shuffle_buffer_size=None,
          prefetch=config.prefetch,
      )
  
      steps_per_epoch = (
          dataset_builder.info.splits["train"].num_examples
          // config.batch_size
      )
  
      if config.num_train_steps <= 0:
          num_steps = int(steps_per_epoch * config.num_epochs)
      else:
          num_steps = config.num_train_steps
  
      if config.steps_per_eval == -1:
          num_validation_examples = dataset_builder.info.splits[
              "validation"
          ].num_examples
          steps_per_eval = num_validation_examples // config.batch_size
      else:
          steps_per_eval = config.steps_per_eval
  
      steps_per_checkpoint = steps_per_epoch * 10
  
      base_learning_rate = (
          config.learning_rate * config.batch_size / 256.0
      )
  
      model_cls = getattr(models, config.model)
      model = create_model(
          model_cls=model_cls, half_precision=config.half_precision
      )
  
      learning_rate_fn = create_learning_rate_fn(
          config, base_learning_rate, steps_per_epoch
      )
  
      state = create_train_state(
          rng, config, model, image_size, learning_rate_fn
      )
      state = restore_checkpoint(state, workdir)
      # step_offset > 0 if restarting from checkpoint
      step_offset = int(state.step)
      state = jax_utils.replicate(state)
  
      p_train_step = jax.pmap(
          functools.partial(
              train_step, learning_rate_fn=learning_rate_fn
          ),
          axis_name="batch",
      )
      p_eval_step = jax.pmap(eval_step, axis_name="batch")
  
      train_metrics = []
      hooks = []
      if jax.process_index() == 0:
          hooks += [
              periodic_actions.Profile(
                  num_profile_steps=5, logdir=workdir
              )
          ]
      train_metrics_last_t = time.time()
      logging.info(
          "Initial compilation, this might take some minutes..."
      )
      for step, batch in zip(range(step_offset, num_steps), train_iter):
          state, metrics = p_train_step(state, batch)
          for h in hooks:
              h(step)
          if step == step_offset:
              logging.info("Initial compilation completed.")
  
          if config.get("log_every_steps"):
              train_metrics.append(metrics)
              if (step + 1) % config.log_every_steps == 0:
                  train_metrics = common_utils.get_metrics(
                      train_metrics
                  )
                  summary = {
                      f"train_{k}": v
                      for k, v in jax.tree_util.tree_map(
                          lambda x: x.mean(), train_metrics
                      ).items()
                  }
                  summary[
                      "steps_per_second"
                  ] = config.log_every_steps / (
                      time.time() - train_metrics_last_t
                  )
                  writer.write_scalars(step + 1, summary)
                  train_metrics = []
                  train_metrics_last_t = time.time()
  
          if (step + 1) % steps_per_epoch == 0:
              epoch = step // steps_per_epoch
              eval_metrics = []
  
              # sync batch statistics across replicas
              state = sync_batch_stats(state)
              for _ in range(steps_per_eval):
                  eval_batch = next(eval_iter)
                  metrics = p_eval_step(state, eval_batch)
                  eval_metrics.append(metrics)
              eval_metrics = common_utils.get_metrics(eval_metrics)
              summary = jax.tree_util.tree_map(
                  lambda x: x.mean(), eval_metrics
              )
              logging.info(
                  "eval epoch: %d, loss: %.4f, accuracy: %.2f",
                  epoch,
                  summary["loss"],
                  summary["accuracy"] * 100,
              )
              writer.write_scalars(
                  step + 1,
                  {f"eval_{key}": val for key, val in summary.items()},
              )
              writer.flush()
          if (
              step + 1
          ) % steps_per_checkpoint == 0 or step + 1 == num_steps:
              state = sync_batch_stats(state)
              save_checkpoint(state, workdir)
  
      # Wait until computations are done before exiting
      jax.random.normal(jax.random.key(0), ()).block_until_ready()
  
      return state
  
imports: import functools
  import time
  from typing import Any
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from flax import jax_utils
  from flax.training import checkpoints
  from flax.training import common_utils
  from flax.training import dynamic_scale
  from flax.training import train_state
  import jax
  from jax import lax
  import jax.numpy
  from jax import random
  import ml_collections
  import optax
  import tensorflow
  import tensorflow_datasets
  import input_pipeline
  import models
label_desc: Use image node_id 14 for calling a global function with example usage: train_and_evaluate(config, workdir) and returns: state