code: def restore_pretrained_big_vision_checkpoint(
      checkpoint_path: str,
  ) -> train_utils.TrainState:
      """Loads and converts a big_vision checkpoint to a scenic train state.
  
      The model weights, global step and accumulated train time are extracted.
      Optimizer state, such as the momentum, is not extracted.
  
      Args:
        checkpoint_path: Path to big_vision checkpoint.
  
      Returns:
        restored_train_state: Scenic train state with model weights, global step
          and accumulated training time.
      """
  
      def unflatten_dict(
          flattened: Dict[str, Any],
          separator: str = "/",
          leaf_idx: int = -1,
      ) -> Dict[str, Any]:
          unflattened = {}
          for k, v in flattened.items():
              subtree = unflattened
              if leaf_idx != 0:
                  path = k.split(separator)[:leaf_idx]
              else:
                  path = k.split(separator)
              for k2 in path[:-1]:
                  if k2 not in subtree:
                      subtree[k2] = {}
                  subtree = subtree[k2]
              subtree[path[-1]] = v
          return unflattened
  
      logging.info(
          "Loading big_vision checkpoint from %s", checkpoint_path
      )
      checkpoint_data = np.load(gfile.GFile(checkpoint_path, "rb"))
      tree = unflatten_dict(checkpoint_data, separator="/", leaf_idx=0)
  
      restored_params = tree["opt"]["target"]
      restored_params = checkpoints.convert_pre_linen(restored_params)
      restored_params = dict(restored_params)
      train_state = train_utils.TrainState()
      # pytype: disable=wrong-arg-types
      restored_train_state = (
          train_state.replace(  # pytype: disable=attribute-error
              params=restored_params
          )
      )
      # pytype: enable=wrong-arg-types
  
      return restored_train_state
  
file_path: scenic/scenic/projects/polyvit/train_utils.py
imports: from typing import Any, Dict, Optional, Tuple
  from absl import logging
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.dataset_lib import datasets
  from scenic.train_lib import train_utils
  from tensorflow.io import gfile
label_desc: Use image node_id 4 for calling a global function with example usage: restore_pretrained_big_vision_checkpoint(checkpoint_path) and returns: restored_train_state, unflattened