code: def train_and_evaluate(
      config: ml_collections.ConfigDict, workdir: str
  ) -> train_state.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:
        The train state (which includes the `.params`).
      """
      # We only support single-host training.
      assert jax.process_count() == 1
  
      # Create writer for logs.
      writer = metric_writers.create_default_writer(workdir)
      writer.write_hparams(config.to_dict())
  
      # Get datasets, organized by split.
      logging.info("Obtaining datasets.")
      datasets = input_pipeline.get_datasets(
          config.batch_size,
          add_virtual_node=config.add_virtual_node,
          add_undirected_edges=config.add_undirected_edges,
          add_self_loops=config.add_self_loops,
      )
      train_iter = iter(datasets["train"])
  
      # Create and initialize the network.
      logging.info("Initializing network.")
      rng = jax.random.key(0)
      rng, init_rng = jax.random.split(rng)
      init_graphs = next(datasets["train"].as_numpy_iterator())
      init_graphs = replace_globals(init_graphs)
      init_net = create_model(config, deterministic=True)
      params = jax.jit(init_net.init)(init_rng, init_graphs)
      parameter_overview.log_parameter_overview(params)
  
      # Create the optimizer.
      tx = create_optimizer(config)
  
      # Create the training state.
      net = create_model(config, deterministic=False)
      state = train_state.TrainState.create(
          apply_fn=net.apply, params=params, tx=tx
      )
  
      # Set up checkpointing of the model.
      # The input pipeline cannot be checkpointed in its current form,
      # due to the use of stateful operations.
      checkpoint_dir = os.path.join(workdir, "checkpoints")
      ckpt = checkpoint.Checkpoint(checkpoint_dir, max_to_keep=2)
      state = ckpt.restore_or_initialize(state)
      initial_step = int(state.step) + 1
  
      # Create the evaluation state, corresponding to a deterministic model.
      eval_net = create_model(config, deterministic=True)
      eval_state = state.replace(apply_fn=eval_net.apply)
  
      # Hooks called periodically during training.
      report_progress = periodic_actions.ReportProgress(
          num_train_steps=config.num_train_steps, writer=writer
      )
      profiler = periodic_actions.Profile(
          num_profile_steps=5, logdir=workdir
      )
      hooks = [report_progress, profiler]
  
      # Begin training loop.
      logging.info("Starting training.")
      train_metrics = None
      for step in range(initial_step, config.num_train_steps + 1):
          # Split PRNG key, to ensure different 'randomness' for every step.
          rng, dropout_rng = jax.random.split(rng)
  
          # Perform one step of training.
          with jax.profiler.StepTraceAnnotation("train", step_num=step):
              graphs = jax.tree_util.tree_map(
                  np.asarray, next(train_iter)
              )
              state, metrics_update = train_step(
                  state, graphs, rngs={"dropout": dropout_rng}
              )
  
              # Update metrics.
              if train_metrics is None:
                  train_metrics = metrics_update
              else:
                  train_metrics = train_metrics.merge(metrics_update)
  
          # Quick indication that training is happening.
          logging.log_first_n(
              logging.INFO, "Finished training step %d.", 10, step
          )
          for hook in hooks:
              hook(step)
  
          # Log, if required.
          is_last_step = step == config.num_train_steps - 1
          if step % config.log_every_steps == 0 or is_last_step:
              writer.write_scalars(
                  step,
                  add_prefix_to_keys(train_metrics.compute(), "train"),
              )
              train_metrics = None
  
          # Evaluate on validation and test splits, if required.
          if step % config.eval_every_steps == 0 or is_last_step:
              eval_state = eval_state.replace(params=state.params)
  
              splits = ["validation", "test"]
              with report_progress.timed("eval"):
                  eval_metrics = evaluate_model(
                      eval_state, datasets, splits=splits
                  )
              for split in splits:
                  writer.write_scalars(
                      step,
                      add_prefix_to_keys(
                          eval_metrics[split].compute(), split
                      ),
                  )
  
          # Checkpoint model, if required.
          if step % config.checkpoint_every_steps == 0 or is_last_step:
              with report_progress.timed("checkpoint"):
                  ckpt.save(state)
  
      return state
  
imports: import os
  from typing import Any, Dict, Iterable, Tuple, Optional
  from absl import logging
  from clu import checkpoint
  from clu import metric_writers
  from clu import metrics
  from clu import parameter_overview
  from clu import periodic_actions
  import flax
  import flax.core
  import flax.linen
  from flax.training import train_state
  import jax
  import jax.numpy
  import jraph
  import ml_collections
  import numpy
  import optax
  import sklearn.metrics
  import tensorflow
  import input_pipeline
  import models
label_desc: Use image node_id 12 for calling a global function with example usage: train_and_evaluate(config, workdir) and returns: state