def predict_and_evaluate(config, workdir, ckpt_path=None):
  """Runs a testing 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.
    ckpt_path: The checkpoint to evaluate. If not specified, use the latest
      checkpoint.
  """
  logging.info('Starting testing at %s', workdir)
  tf.io.gfile.makedirs(workdir)

  rng = jax.random.PRNGKey(config.seed)
  # Build input pipeline.
  rng, data_rng = jax.random.split(rng)
  data_rng = jax.random.fold_in(data_rng, jax.process_index())
  test_ds = []
  for split in config.dataset.test_splits:
    ds = input_pipeline.create_val_dataset(
        config.dataset, split, config.dataset.test_per_device_batch_size,
        config.dataset.test_pad_last_batch)
    test_ds.append(ds)

  # Initialize model.
  inputs = train_utils.get_init_inputs(test_ds[0])
  rng, model_rng = jax.random.split(rng)
  predict_config = models.TransformerConfig(**config.model.to_dict())
  predict_config = predict_config.replace(decode=True)
  model = models.Model(predict_config)
  state = train_utils.create_train_state(
      model, config, model_rng, inputs=inputs)

  writer = metric_writers.create_default_writer(
      workdir, just_logging=jax.process_index() > 0)

  # Set up checkpointing of the model and the input pipeline.
  checkpoint_dir = os.path.join(workdir, 'checkpoints')
  ckpt = checkpoint.MultihostCheckpoint(checkpoint_dir, max_to_keep=3)

  logging.info('Testing and evaluating checkpoint %s', ckpt_path)
  try:
    state = ckpt.restore(state, ckpt_path)
  except FileNotFoundError:
    state = ckpt.restore_or_initialize(state)
  step = int(state.step)

  p_pred_step = jax.pmap(
      functools.partial(predict_step, config=predict_config),
      axis_name='batch',
      static_broadcasted_argnums=(3,))
  p_init_cache = jax.pmap(
      functools.partial(init_cache, config=predict_config), axis_name='batch')

  # Distribute testing.
  state = flax_utils.replicate(state)
  with metric_writers.ensure_flushes(writer):
    test_metrics = {}
    for ds, split in zip(test_ds, config.dataset.test_splits):
      ds_metrics = evaluate_sequence_accuracy(p_pred_step, p_init_cache, state,
                                              ds, config, split, workdir,
                                              config.num_test_steps)
      ds_metrics = {f'{k}_{split}': v for k, v in ds_metrics.items()}
      test_metrics.update(ds_metrics)
    writer.write_scalars(step, test_metrics)
Runs a testing 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.
    ckpt_path: The checkpoint to evaluate. If not specified, use the latest
      checkpoint.
  
language/gscan/xattn_model/predict.py