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)

  # Build the flax_model and the optimizers.
  _, init_rng = jax.random.split(rng)
  input_spec = [(
      dataset.meta_data['input_shape'],  # bs, t, h, w, 3
      dataset.meta_data.get('input_dtype', dtype))]
  model_dict, train_state, _ = create_train_state(input_spec, config, init_rng,
                                                  False)
  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)

  # Eval step pmap.
  eval_step_pmapped = jax.pmap(
      functools.partial(
          eval_step,
          model_dict=model_dict,
          config=config,
          metric_params=metric_params,
      ),
      axis_name='device',
      # We can donate the buffer of batch.
      donate_argnums=(1),
  )

  log_sample_size = config.logging.get('log_sample_size', 2)
  label_names, result_dir, write_executor = None, None, None
  if config.eval.get('results_dir') is not None:
    label_names = eval_utils.load_label_names(config)
    result_dir = eval_utils.get_result_dir(config)
    write_executor = concurrent.futures.ThreadPoolExecutor(100)

  all_metrics, batch_samples = {}, 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

    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)
    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)
    eval_metrics = {}
    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 = f"{data_split}{config.eval.get('suffix', '')}"

      eval_features, batch_outputs = [], {}
      with metric_writers.ensure_flushes(writer):
        for step in range(num_steps):
          eval_batch = next(data_iter)
          if config.dataset_configs.get('one_hot_labels', False):
            eval_batch.pop('label', None)
          batch_features, batch_outputs = eval_step_pmapped(
              train_state_replicated, eval_batch)
          batch_features = jax.device_get(jax_utils.unreplicate(batch_features))
          eval_features.append(batch_features)
          # Quick indication that evaluation is happening.
          logging.log_first_n(logging.INFO, 'Finished evaluation step %d.', 5,
                              step)
          for h in hooks:
            try:
              h(start_step + step)
            except:  # pylint: disable=bare-except
              logging.exception('Hook failed')
              continue
          if write_executor is not None or config.eval.get(
              'enable_lpips', False):
            batch_outputs = jax.device_get(batch_outputs)
            if config.eval.get('enable_lpips', False):
              # batch_outputs has been all_gathered
              batch_samples = jax_utils.unreplicate(batch_outputs)
              lpips_scores = _compute_lpips_one_step(batch_samples,
                                                     metric_params)
              eval_features[-1].update(lpips_scores)
              # only keep the local batch for writing
              batch_samples = eval_utils.get_local_batch(batch_samples)
            else:
              batch_samples = jax.tree_util.tree_map(
                  lambda x: x.reshape(-1, *x.shape[2:]), batch_outputs)
            if write_executor is not None:
              # batch_samples is local flattened batch as np.array
              # batch_features is global flattened batch as np.array
              batch_features = eval_utils.get_local_batch(batch_features)
              batch_samples.update(batch_features)
              eval_utils.write_batch_samples(
                  batch_samples, step, write_executor, result_dir,
                  f'{metric_suffix}_step{global_step}', label_names, config)

        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,
            model_suffix_list=['', '_ema'])
        eval_metrics = {
            f'eval_{metric_suffix}/{k}': v for k, v in eval_metrics.items()
        }
        writer.write_scalars(global_step, eval_metrics)
        all_metrics.update(eval_metrics)
        # Sample from last batch.
        batch_samples = jax.device_get(
            jax.tree_util.tree_map(lambda x: x[0, :log_sample_size],
                                   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)

  if write_executor is not None:
    write_executor.shutdown()
  return train_state, all_metrics, batch_samples
Main evaluation loop lives in this function.
videogvt/trainers/vqgan_trainer.py