code: def restore_pretrained_checkpoint(
      checkpoint_path: str,
      train_state: Optional[train_utils.TrainState] = None,
      assert_exist: bool = False,
      step: Optional[int] = None,
  ) -> train_utils.TrainState:
      """Restores the last checkpoint.
  
      First restores the checkpoint, which is an instance of TrainState that holds
      the state of training. This function also take care converting pre-Linen
      checkpoints.
  
      Args:
        checkpoint_path: Directory for saving 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}"
              )
      restored_train_state = checkpoints.restore_checkpoint(
          checkpoint_path, None, step
      )
      if restored_train_state is None:
          raise ValueError(
              "No checkpoint for the pretrained model is found in: "
              f"{checkpoint_path}"
          )
      if "params" in restored_train_state:
          # restored_train_state was trained using optax
          restored_params = flax.core.freeze(
              restored_train_state["params"]
          )
      else:
          # restored_train_state was trained using flax.optim. Note that this does
          # not convert the naming of pre-Linen checkpoints.
          restored_params = restored_train_state["optimizer"]["target"]
          if "params" in restored_params:  # Backward compatibility.
              restored_params = restored_params["params"]
              restored_params = dict(
                  checkpoints.convert_pre_linen(restored_params)
              )
          restored_params = flax.core.freeze(restored_params)
  
      restored_model_state = flax.core.freeze(
          restored_train_state["model_state"]
      )
  
      if not train_state:
          train_state = train_utils.TrainState()
          params = restored_params
      else:
          # Inspect and compare the parameters of the model with the init-model.
          params = inspect_params(
              expected_params=train_state.params,
              restored_params=restored_params,
              fail_if_extra=False,
              fail_if_missing=False,
              fail_if_shapes_mismatch=False,
          )
      train_state = train_state.replace(
          # Inspect and compare the parameters of the model with the init-model.
          params=params,
          model_state=restored_model_state,
          global_step=int(restored_train_state["global_step"]),
          rng=restored_train_state["rng"],
          metadata=restored_train_state.get("metadata", None),
      )
      return train_state
  
file_path: scenic/scenic/train_lib/pretrain_utils.py
imports: from collections import abc
  import os
  import re
  from typing import Any, Dict, Mapping, List, Optional, Union
  from absl import logging
  import flax
  from flax.training import checkpoints
  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_pretrained_checkpoint(checkpoint_path, train_state, assert_exist, step) and returns: train_state