code: def convert_big_vision_to_scenic_checkpoint(
      checkpoint_path: str,
      train_state: Optional[train_utils.TrainState] = None,
  ) -> 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.
  
      Returns:
        restored_train_state: Scenic train state with model weights, global step
          and accumulated training time.
      """
  
      def _recover_tree(keys, values):
          tree = {}
          sub_trees = collections.defaultdict(list)
          for k, v in zip(keys, values):
              if "/" not in k:
                  tree[k] = v
              else:
                  k_left, k_right = k.split("/", 1)
                  sub_trees[k_left].append((k_right, v))
          for k, kv_pairs in sub_trees.items():
              k_subtree, v_subtree = zip(*kv_pairs)
              tree[k] = _recover_tree(k_subtree, v_subtree)
          return tree
  
      logging.info(
          "Loading big_vision checkpoint from %s", checkpoint_path
      )
      checkpoint_npz = np.load(gfile.GFile(checkpoint_path, "rb"))
      keys, values = zip(*list(checkpoint_npz.items()))
      checkpoint = _recover_tree(keys, values)
      restored_params = checkpoints.convert_pre_linen(
          checkpoint.get("params", checkpoint)
      )
      if train_state:
          restored_params = pretrain_utils.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()
  
      global_step = None
      if "opt" in checkpoint:
          global_step = scenic_optax.get_step(checkpoint["opt"])
  
      # pytype: disable=wrong-arg-types
      restored_train_state = (
          train_state.replace(  # pytype: disable=attribute-error
              global_step=global_step,
              params=restored_params,
          )
      )
      # pytype: enable=wrong-arg-types
  
      return restored_train_state
  
imports: import collections
  import functools
  from typing import Any, Callable, Dict, Iterator, Tuple, Optional, Type
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from clu import platform
  from flax import jax_utils
  import flax.linen
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  import optax
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import base_model
  from scenic.train_lib import optax
  from scenic.train_lib import pretrain_utils
  from scenic.train_lib import train_utils
  from scenic.train_lib.transfer import fewshot_utils
  from tensorflow.io import gfile
label_desc: Use image node_id 1 for calling a global function with example usage: convert_big_vision_to_scenic_checkpoint(checkpoint_path, train_state) and returns: restored_train_state, tree