def evaluate(*, rng: jnp.ndarray, config: ml_collections.ConfigDict,
             dataset: dataset_utils.Dataset, workdir: str,
             writer: metric_writers.MetricWriter):
  """Main evaluation loop lives in this function."""
  lead_host = jax.process_index() == 0
  dtype = train_utils.get_dtype(config)
  config, ckpt_manager, ckpt_list = eval_utils.get_eval_jobs(workdir, config)

  tokenizer_dict = maskgvt_train_lib.load_vq_tokenizer(config)
  writer.write_hparams(eval_utils.flatten_config(config))

  # Build the flax_model and the optimizers.
  rng, init_rng = jax.random.split(rng, 2)
  input_spec = [(
      dataset.meta_data['input_shape'],  # bs, t, h, w, 3
      dataset.meta_data.get('input_dtype', dtype))]

  # Get evaluation task
  task_reg = task_registry.get_task_registry(config, False)
  task_ids = config.get('tasks', ('full_generation',))
  if 'task' in config.eval:
    task_ids = (config.eval.task,)
  assert len(task_ids) == 1
  # Loop through all tasks during parallel evaluation
  task_ids_iter = itertools.cycle(task_ids)

  model, model_config, train_state, _ = create_train_state(
      input_spec, task_reg[task_ids[0]], tokenizer_dict['tokenizer'],
      config, init_rng)
  del model
  # to avoid loading cache from checkpoints
  pred_model_config = model_config.replace(decode=True)
  pred_model = transformer_lm.TransformerLM(pred_model_config)

  total_train_steps, _ = scenic_train_utils.get_num_training_steps(
      config, dataset.meta_data)

  # Load metric params
  metric_params = eval_utils.load_metric_params(config)

  # Pmap.
  eval_step_pmapped = jax.pmap(
      functools.partial(
          eval_step,
          flax_model=pred_model,
          tokenizer_dict=tokenizer_dict,
          config=config,
          metric_params=metric_params,
      ),
      static_broadcasted_argnums=(3,),
      axis_name='device',
      # We can donate the buffer of batch.
      donate_argnums=(1),
  )

  init_cache_pmapped = jax.pmap(
      functools.partial(
          initialize_cache,
          task=task_reg[task_ids[0]],
          model=pred_model,
          input_spec=input_spec,
          tokenizer_fn=tokenizer_dict['tokenizer'],
          config=config),
      axis_name='device'
  )

  log_sample_size = config.logging.get('log_sample_size', 2)

  eval_metrics, batch_samples, batch_sample_list = None, None, []
  for ckpt_path in ckpt_list:
    # Restores the model
    if not gfile.exists(ckpt_path):
      logging.warn(
          'Unable to evaluate ckpt %s because it does not exist. '
          'If this is a parallel evaluation job, try to increase '
          'config.logging.checkpoint_kept or use more accelerators.', ckpt_path)
      continue
    task = task_reg[next(task_ids_iter)]
    train_state = train_utils.restore_checkpoint(
        ckpt_path,
        train_state,
        is_legacy_checkpoint=config.eval_from.get('legacy_checkpoint', False))
    global_step = int(train_state.global_step)
    ckpt_rng = jax.random.fold_in(rng, global_step)
    start_step = global_step  # Step offset for logging
    train_state_replicated = jax_utils.replicate(train_state)
    logging.info('Starting evaluation at global step %d', global_step)

    is_final_eval = False
    if global_step == total_train_steps or config.get('eval_only', False):
      is_final_eval = True
    data_iters = eval_utils.get_eval_data_iters(dataset, config, is_final_eval)
    for data_split, (data_iter, num_steps, total_num_examples,
                     num_repeats) in data_iters.items():
      if not config.get('eval_only', False):
        report_progress = eval_utils.ReportEvaluationProgress(
            num_train_steps=None, writer=writer)
      else:
        report_progress = periodic_actions.ReportProgress(
            num_train_steps=start_step + num_steps, writer=writer)
      hooks = [report_progress]
      if config.get('xprof', False) and lead_host:
        hooks.append(
            periodic_actions.Profile(num_profile_steps=5, logdir=workdir))

      logging.info(
          'Evaluating for %d steps on %s split', num_steps, data_split)
      metric_suffix = data_split + config.eval.get('suffix', '')

      eval_features = []
      with metric_writers.ensure_flushes(writer):
        for step in range(num_steps):
          eval_batch = next(data_iter)
          ckpt_rng, step_rng = jax.random.split(ckpt_rng, 2)
          cache = init_cache_pmapped(eval_batch['inputs'])
          step_rng = jax.random.split(step_rng, jax.local_device_count())
          batch_features, batch_outputs = eval_step_pmapped(
              train_state_replicated, eval_batch, cache, task, step_rng)
          batch_features = jax.device_get(jax_utils.unreplicate(batch_features))
          eval_features.append(batch_features)
          # Quick indication that training is happening.
          logging.log_first_n(logging.INFO, 'Finished evaluation step %d.', 5,
                              step)
          for h in hooks:
            h(start_step + step)
          if config.eval.get('return_samples', False):
            batch_samples = jax.device_get(jax_utils.unreplicate(batch_outputs))
            batch_sample_list.append(batch_samples)
        start_step += num_steps
        eval_features = eval_utils.gather_outputs_with_mask(
            eval_features, num_samples=total_num_examples)
        eval_metrics = eval_utils.compute_eval_metrics(eval_features, config,
                                                       total_num_examples,
                                                       num_repeats)
        eval_metrics = {
            f'eval_{metric_suffix}/{k}': v for k, v in eval_metrics.items()
        }
        writer.write_scalars(global_step, eval_metrics)
        # Sample from last batch.
        batch_samples = jax.device_get(jax_utils.unreplicate(batch_outputs))
        log_samples(batch_samples, writer, f'eval_{metric_suffix}', global_step,
                    log_sample_size,
                    config.get('dataset_type', 'video') == 'video')

    del train_state_replicated
    ckpt_manager.add_eval_result(ckpt_path, eval_metrics, 0)
    logging.info('Finishing evaluation at global step %d', global_step)

  # Wait until computations are done before exiting.
  jax.random.normal(jax.random.PRNGKey(0), ()).block_until_ready()
  # Return the train and eval summary after last step for regression testing.
  if config.eval.get('return_samples', False):
    samples = jax.tree_util.tree_map(lambda *x: np.concatenate(x),
                                     *batch_sample_list)
  else:
    samples = batch_samples
  return train_state, eval_metrics, samples
Main evaluation loop lives in this function.
videogvt/trainers/lmgvt_trainer.py