code: def restore_checkpoint(
      checkpoint_path: str,
      train_state: Optional[train_utils.TrainState] = None,
      assert_exist: bool = False,
      step: Optional[int] = None,
  ) -> Tuple[train_utils.TrainState, int]:
      """Restores the last checkpoint.
  
      First restores the checkpoint, which is an instance of TrainState that holds
      the state of training.
  
      Args:
        checkpoint_path: Directory to restore the checkpoint.
        train_state: An instance of TrainState that holds the state of
          training.
        assert_exist: 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}"
              )
      if train_state is None:
          raise ValueError(
              "Please use `restore_pretrained_checkpoint` for loading"
              "a checkpoint without providing a Scenic TrainState."
          )
      train_state = checkpoints.restore_checkpoint(
          checkpoint_path, train_state, step
      )
      return train_state, int(train_state.global_step)
  
file_path: scenic/scenic/projects/ncr/utils.py
imports: import os
  from typing import Any, Dict, Optional, Tuple
  from flax import jax_utils
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import numpy
  from scenic.train_lib import train_utils
  from tensorflow.io import gfile
label_desc: Use image node_id 3 for calling a global function with example usage: restore_checkpoint(checkpoint_path, train_state, assert_exist, step) and returns: train_state,  int