code: def restore_checkpoint(
      checkpoint_path: str,
      train_state: Optional[train_utils.TrainState] = None,
      assert_exist: bool = False,
      step: int = None,
  ) -> Tuple[train_utils.TrainState, int]:
      """Restores the last checkpoint.
  
      Supports checkpoints saved either with old Scenic (flax.deprecated.nn) or
      current Scenic (flax.Linen). Therefore, this function can be used for
      evaluation of old or current models.
  
      First restores the checkpoint, which is an instance of TrainState that holds
      the state of training, and then replicates it.
  
      Args:
        checkpoint_path: Directory for saving the checkpoint.
        train_state: An instance of TrainState that holds the state of training.
        assert_exist: bool; Assert that there is at least one checkpoint exists in
          the given path.
        step: Step number to load or None to load latest. If specified,
          checkpoint_path must be a directory.
  
      Returns:
        Training state and an int which is the current step.
      """
      if assert_exist:
          glob_path = os.path.join(checkpoint_path, "checkpoint_*")
          if not gfile.glob(glob_path):
              raise ValueError(
                  "No checkpoint for the pretrained model is found in: "
                  f"{checkpoint_path}"
              )
      restored_train_state = checkpoints.restore_checkpoint(
          checkpoint_path, None, step
      )
  
      if restored_train_state:
          (
              restored_params,
              restored_model_state,
          ) = pretrain_utils.get_params_and_model_state_dict(
              restored_train_state
          )
          restored_params = flax.core.freeze(restored_params)
          restored_model_state = flax.core.freeze(restored_model_state)
          train_state = train_state or train_utils.TrainState()
          if train_state.optimizer:
              new_optimizer = train_state.optimizer.replace(
                  target=restored_params
              )
          else:
              new_optimizer = {"target": restored_params}
          train_state = (
              train_state.replace(  # pytype: disable=attribute-error
                  optimizer=new_optimizer,
                  model_state=restored_model_state,
                  global_step=int(restored_train_state["global_step"]),
                  rng=restored_train_state["rng"],
                  accum_train_time=restored_train_state.get(
                      "accum_train_time", 0
                  ),
              )
          )
      else:
          train_state = train_state or train_utils.TrainState()
  
      return train_state, int(train_state.global_step)
  
file_path: scenic/scenic/projects/vivit/evaluation_lib.py
imports: import os
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union
  from absl import logging
  import flax
  from flax.metrics import tensorboard
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import numpy
  from scenic.train_lib_deprecated import pretrain_utils
  from scenic.train_lib_deprecated import train_utils
  from sklearn.metrics import average_precision_score
  from tensorflow.io import gfile
label_desc: Use image node_id 1 for calling a global function with example usage: restore_checkpoint(checkpoint_path, train_state, assert_exist, step) and returns: train_state,  int