code: def train_and_evaluate(config: default.Config, workdir: str):
      """Runs a training and evaluation loop.
  
      Args:
        config: Configuration to use.
        workdir: Working directory for checkpoints and TF summaries. If this
          contains checkpoint training will be resumed from the latest checkpoint.
      """
      tf.io.gfile.makedirs(workdir)
  
      vocab_path = config.vocab_path
      if vocab_path is None:
          vocab_path = os.path.join(workdir, "sentencepiece_model")
          config.vocab_path = vocab_path
      tf.io.gfile.makedirs(os.path.split(vocab_path)[0])
  
      # Load Dataset
      # ---------------------------------------------------------------------------
      logging.info("Initializing dataset.")
      train_ds, eval_ds, _, encoder = input_pipeline.get_datasets(
          n_devices=jax.local_device_count(),
          config=config,
          vocab_path=vocab_path,
      )
  
      train_iter = iter(train_ds)
      vocab_size = int(encoder.vocab_size())
      eos_id = (
          temperature_sampler.EOS_ID
      )  # Default Sentencepiece EOS token.
  
      def decode_tokens(toks):
          valid_toks = toks[: np.argmax(toks == eos_id) + 1].astype(
              np.int32
          )
          return encoder.detokenize(valid_toks).numpy().decode("utf-8")
  
      def encode_strings(strs, max_len):
          tokenized_batch = np.zeros((len(strs), max_len), np.int32)
          for i, s in enumerate(strs):
              toks = encoder.tokenize(s).numpy()
              # Remove EOS token in prompt.
              tokenized_batch[i, : toks.shape[0] - 1] = toks[:-1]
          return tokenized_batch
  
      tokenized_prompts = encode_strings(
          [config.prompts], config.max_predict_length
      )
  
      logging.info("Initializing model, optimizer, and step functions.")
      # Build Model and Optimizer
      # ---------------------------------------------------------------------------
      model_config = models.TransformerConfig(
          vocab_size=vocab_size,
          output_vocab_size=vocab_size,
          logits_via_embedding=config.logits_via_embedding,
          dtype=jnp.bfloat16 if config.use_bfloat16 else jnp.float32,
          emb_dim=config.emb_dim,
          num_heads=config.num_heads,
          num_layers=config.num_layers,
          qkv_dim=config.qkv_dim,
          mlp_dim=config.mlp_dim,
          max_len=max(
              config.max_target_length, config.max_eval_target_length
          ),
          dropout_rate=config.dropout_rate,
          attention_dropout_rate=config.attention_dropout_rate,
          kernel_init=nn.initializers.xavier_uniform(),
          bias_init=nn.initializers.normal(stddev=1e-6),
          axis_rules=config.axis_rules,
      )
  
      # Mesh definition
      devices_array = utils.create_device_mesh(config)
      mesh = Mesh(devices_array, config.mesh_axes)
  
      # print(mesh.shape)
      # exit()
  
      start_step = 0
      rng = jax.random.PRNGKey(config.seed)
      rng, init_rng = jax.random.split(rng)
      rng, inference_rng = random.split(rng)
  
      def constructor(config: models.TransformerConfig, key: jax.Array):
          return models.TransformerLM(config, rngs=nnx.Rngs(params=key))
  
      learning_rate_fn = create_learning_rate_schedule(
          learning_rate=config.learning_rate,
          warmup_steps=config.warmup_steps,
      )
  
      optimizer = optax.adamw(
          learning_rate_fn,
          b1=0.9,
          b2=0.98,
          eps=1e-9,
          weight_decay=config.weight_decay,
      )
  
      state, state_sharding = utils.setup_initial_state(
          constructor, optimizer, model_config, init_rng, mesh
      )
      data_sharding = NamedSharding(mesh, P(config.data_sharding))
  
      if config.restore_checkpoints:
          # Restore unreplicated optimizer + model state from last checkpoint.
          state = checkpoints.restore_checkpoint(workdir, state)
          # Grab last step.
          start_step = int(state.step)
  
      writer = metric_writers.create_default_writer(
          workdir, just_logging=jax.process_index() > 0
      )
      if start_step == 0:
          writer.write_hparams(dataclasses.asdict(config))
  
      # compile multidevice versions of train/eval/predict step fn.
      jit_train_step = jax.jit(
          train_step,
          in_shardings=(
              state_sharding,
              data_sharding,
              None,
          ),  # type: ignore
          out_shardings=(state_sharding, None),  # type: ignore
          static_argnums=(2, 3),
          donate_argnums=0,
      )
  
      jit_eval_step = jax.jit(
          eval_step,
          in_shardings=(
              state_sharding.params,
              data_sharding,
          ),  # type: ignore
          out_shardings=None,  # type: ignore
          static_argnums=(2, 3),
      )
  
      # Since the inputs and rngkey args for predict_step will be batched,
      # we must vmap them, otherwise the global arrays will be seen in each device
      jit_pred_step = jax.jit(
          jax.vmap(
              predict_step,
              in_axes=(
                  0,
                  jax.tree_map(lambda x: None, state.params),
                  0,
                  None,
                  None,
                  None,
                  None,
                  None,
                  None,
              ),
          ),
          in_shardings=(
              data_sharding,
              state_sharding.params,
              data_sharding,
          ),  # type: ignore
          out_shardings=data_sharding,  # type: ignore
          static_argnums=tuple(range(3, 9)),
      )
  
      # Main Train Loop
      # ---------------------------------------------------------------------------
  
      # We init the first set of dropout PRNG keys, but update it afterwards inside
      # the main pmap"d training update for performance.
      dropout_rngs = rng
  
      logging.info("Starting training loop.")
      hooks = []
      report_progress = periodic_actions.ReportProgress(
          num_train_steps=config.num_train_steps, writer=writer
      )
      if jax.process_index() == 0:
          hooks += [
              report_progress,
              periodic_actions.Profile(
                  logdir=workdir, num_profile_steps=5
              ),
          ]
      train_metrics = []
      with metric_writers.ensure_flushes(writer):
          for step in range(start_step, config.num_train_steps):
              is_last_step = step == config.num_train_steps - 1
  
              # Shard data to devices and do a training step.
              with jax.profiler.StepTraceAnnotation(
                  "train", step_num=step
              ):
                  batch = next(train_iter)
                  batch = jax.tree_map(lambda x: jnp.asarray(x), batch)
                  state, metrics = jit_train_step(
                      state, batch, learning_rate_fn, 0.0, dropout_rngs
                  )
                  train_metrics.append(metrics)
  
              # Quick indication that training is happening.
              logging.log_first_n(
                  logging.INFO, "Finished training step %d.", 5, step
              )
              for h in hooks:
                  h(step)
  
              # Periodic metric handling.
              if step % config.eval_every_steps == 0 or is_last_step:
                  with report_progress.timed("training_metrics"):
                      logging.info("Gathering training metrics.")
                      train_metrics = common_utils.stack_forest(
                          train_metrics
                      )
                      lr = train_metrics.pop("learning_rate").mean()
                      metrics_sums = jax.tree_util.tree_map(
                          jnp.sum, train_metrics
                      )
                      denominator = metrics_sums.pop("denominator")
                      summary = jax.tree_util.tree_map(
                          lambda x: x / denominator, metrics_sums
                      )  # pylint: disable=cell-var-from-loop
                      summary["learning_rate"] = lr
                      summary["perplexity"] = jnp.clip(
                          jnp.exp(summary["loss"]), a_max=1.0e4
                      )
                      summary = {
                          "train_" + k: v for k, v in summary.items()
                      }
                      writer.write_scalars(step, summary)
                      train_metrics = []
  
                  with report_progress.timed("eval"):
                      eval_results = evaluate(
                          jit_eval_step=jit_eval_step,
                          state=state,
                          eval_ds=eval_ds,
                          num_eval_steps=config.num_eval_steps,
                      )
                      # (clipped) perplexity after averaging log-perplexitie
                      eval_results["perplexity"] = jnp.clip(
                          jnp.exp(eval_results["loss"]), a_max=1.0e4
                      )
                      writer.write_scalars(
                          step,
                          {
                              "eval_" + k: v
                              for k, v in eval_results.items()
                          },
                      )
  
                  with report_progress.timed("generate_text"):
                      exemplars = generate_prediction(
                          jit_pred_step=jit_pred_step,
                          static=state.graphdef,
                          params=state.params,
                          tokenized_prompts=tokenized_prompts,
                          eos_id=eos_id,
                          inference_rng=inference_rng,
                          decode_tokens=decode_tokens,
                          config=config,
                          model_config=model_config,
                      )
                      writer.write_texts(step, {"samples": exemplars})
  
              # Save a checkpoint on one host after every checkpoint_freq steps.
              save_checkpoint = (
                  step % config.checkpoint_every_steps == 0
                  or is_last_step
              )
              if config.save_checkpoints and save_checkpoint:
                  logging.info("Saving checkpoint step %d.", step)
                  with report_progress.timed("checkpoint"):
                      checkpoints.save_checkpoint_multiprocess(
                          workdir, state, step
                      )
  
imports: import collections
  import dataclasses
  import os
  import input_pipeline
  import jax
  import jax.numpy
  import models
  import numpy
  import optax
  import temperature_sampler
  import tensorflow
  import utils
  from absl import logging
  from clu import metric_writers, periodic_actions
  from configs import default
  from jax import random
  from jax.sharding import Mesh, NamedSharding
  from jax.sharding import PartitionSpec
  from utils import HasCache, TrainState
  from flax import linen
  from flax.experimental import nnx
  from flax.training import checkpoints, common_utils
label_desc: Use image node_id 14 for calling a global function with example usage: train_and_evaluate(config, workdir) and returns: encoder, tokenized_batch, models