code: def train_and_evaluate(
      config: ml_collections.ConfigDict, 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,
          predict_ds,
          encoder,
      ) = input_pipeline.get_wmt_datasets(
          n_devices=jax.local_device_count(),
          config=config,
          reverse_translation=config.reverse_translation,
          vocab_path=vocab_path,
      )
  
      train_iter = iter(train_ds)
      vocab_size = int(encoder.vocab_size())
      eos_id = decode.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")
  
      if config.num_predict_steps > 0:
          predict_ds = predict_ds.take(config.num_predict_steps)
  
      logging.info("Initializing model, optimizer, and step functions.")
  
      dtype = preferred_dtype(config)
  
      # Build Model and Optimizer
      # ---------------------------------------------------------------------------
      train_config = models.TransformerConfig(
          vocab_size=vocab_size,
          output_vocab_size=vocab_size,
          share_embeddings=config.share_embeddings,
          logits_via_embedding=config.logits_via_embedding,
          dtype=dtype,
          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,
          deterministic=False,
          decode=False,
          kernel_init=nn.initializers.xavier_uniform(),
          bias_init=nn.initializers.normal(stddev=1e-6),
      )
      eval_config = train_config.replace(deterministic=True)
      predict_config = train_config.replace(
          deterministic=True, decode=True
      )
  
      start_step = 0
      rng = jax.random.key(config.seed)
      rng, init_rng = jax.random.split(rng)
      input_shape = (
          config.per_device_batch_size,
          config.max_target_length,
      )
      target_shape = (
          config.per_device_batch_size,
          config.max_target_length,
      )
  
      m = models.Transformer(eval_config)
      initial_variables = jax.jit(m.init)(
          init_rng,
          jnp.ones(input_shape, jnp.float32),
          jnp.ones(target_shape, jnp.float32),
      )
  
      # Create train state with Adam optimizer and weight decay.
      learning_rate_fn = create_learning_rate_schedule(
          learning_rate=config.learning_rate,
          warmup_steps=config.warmup_steps,
      )
      dynamic_scale = None
      if dtype == jnp.float16:
          dynamic_scale = dynamic_scale_lib.DynamicScale()
      state = TrainState.create(
          apply_fn=m.apply,
          params=initial_variables["params"],
          tx=optax.adamw(
              learning_rate=learning_rate_fn,
              b1=0.9,
              b2=0.98,
              eps=1e-9,
              weight_decay=config.weight_decay,
          ),
          dynamic_scale=dynamic_scale,
      )
  
      # We access model params only via state.params
      del initial_variables
  
      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(dict(config))
  
      # Replicate state.
      state = jax_utils.replicate(state)
  
      # compile multidevice versions of train/eval/predict step and cache init fn.
      p_train_step = jax.pmap(
          functools.partial(
              train_step,
              config=train_config,
              learning_rate_fn=learning_rate_fn,
              label_smoothing=config.label_smoothing,
          ),
          axis_name="batch",
          donate_argnums=(0,),
      )  # pytype: disable=wrong-arg-types
      p_eval_step = jax.pmap(
          functools.partial(eval_step, config=eval_config),
          axis_name="batch",
      )
      p_init_cache = jax.pmap(
          functools.partial(
              initialize_cache,
              max_decode_len=config.max_predict_length,
              config=predict_config,
          ),
          axis_name="batch",
      )
      p_pred_step = jax.pmap(
          functools.partial(
              predict_step,
              config=predict_config,
              beam_size=config.beam_size,
          ),
          axis_name="batch",
          static_broadcasted_argnums=(3, 4),
      )  # eos token, max_length are constant
  
      # 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 = jax.random.split(rng, jax.local_device_count())
      del 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 = common_utils.shard(
                      jax.tree_util.tree_map(
                          np.asarray, next(train_iter)
                      )
                  )
                  state, metrics = p_train_step(
                      state, batch, dropout_rng=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.get_metrics(
                          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 = {
                          "train_" + k: v for k, v in summary.items()
                      }
                      writer.write_scalars(step, summary)
                      train_metrics = []
  
                  with report_progress.timed("eval"):
                      eval_results = evaluate(
                          p_eval_step=p_eval_step,
                          params=state.params,
                          eval_ds=eval_ds,
                          num_eval_steps=config.num_eval_steps,
                      )
                      writer.write_scalars(
                          step,
                          {
                              "eval_" + k: v
                              for k, v in eval_results.items()
                          },
                      )
  
                  with report_progress.timed("translate_and_bleu"):
                      (
                          exemplars,
                          bleu_score,
                      ) = translate_and_calculate_bleu(
                          p_pred_step=p_pred_step,
                          p_init_cache=p_init_cache,
                          params=state.params,
                          predict_ds=predict_ds,
                          decode_tokens=decode_tokens,
                          max_predict_length=config.max_predict_length,
                      )
                      writer.write_scalars(step, {"bleu": bleu_score})
                      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, jax_utils.unreplicate(state), step
                      )
  
imports: import collections
  import functools
  import os
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from flax import jax_utils
  from flax import linen
  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
  import jax.numpy
  import ml_collections
  import numpy
  import optax
  import tensorflow
  import bleu
  import decode
  import input_pipeline
  import models
label_desc: Use image node_id 16 for calling a global function with example usage: train_and_evaluate(config, workdir) and returns: encoder