code: def convert_big_vision_to_scenic_checkpoint(
      checkpoint_path: str,
      train_state: Optional[train_utils.TrainState] = None,
      convert_to_linen: bool = True,
  ) -> train_utils.TrainState:
      """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.
        train_state: A Scenic TrainState object.
        convert_to_linen: Whether to convert to Linen format.
  
      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"]
      if convert_to_linen:
          restored_params = checkpoints.convert_pre_linen(
              restored_params
          )
      restored_params = dict(restored_params)
      if train_state:
          restored_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,
          )
      else:
          train_state = train_utils.TrainState()
  
      # pytype: disable=wrong-arg-types
      restored_train_state = (
          train_state.replace(  # pytype: disable=attribute-error
              global_step=int(tree["opt"]["state"]["step"]),
              params=restored_params,
          )
      )
      # pytype: enable=wrong-arg-types
  
      return restored_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 5 for calling a global function with example usage: convert_big_vision_to_scenic_checkpoint(checkpoint_path, train_state, convert_to_linen) and returns: restored_train_state, unflattened