code: def restore_checkpoint(
      checkpoint_path: str,
      train_state: Optional[TrainState] = None,
      assert_exist: bool = False,
      step: Optional[int] = None,
  ) -> Tuple[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 or filename to restore the checkpoint from.
        train_state: An instance of TrainState that holds the state of training.
        assert_exist: Assert that there is at least one checkpoint 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:
          if "checkpoint_" in checkpoint_path.split("/")[-1]:
              glob_path = checkpoint_path
          else:
              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)
  
imports: import collections.abc
  import copy
  import functools
  import os
  import re
  import time
  from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple, Union
  from absl import logging
  from clu import metric_writers
  import flax
  from flax import jax_utils
  from flax import struct
  import flax.linen
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  import optax
  from scenic.common_lib import debug_utils
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  from scenic.train_lib import optimizers
  from tensorflow.io import gfile
label_desc: Use image node_id 10 for calling a global function with example usage: restore_checkpoint(checkpoint_path, train_state, assert_exist, step) and returns: train_state,  int