code: def init_vit_from_train_state(
      train_state: Any,
      restored_train_state: Any,
      model_cfg: ml_collections.ConfigDict,
      restored_model_cfg: ml_collections.ConfigDict,
  ) -> Any:
      """Updates the init_params with data from restored_params."""
      del restored_model_cfg
      params_dict = flax.traverse_util.flatten_dict(
          flax.core.unfreeze(train_state.params), sep="/"
      )
      restored_params_dict = flax.traverse_util.flatten_dict(
          flax.core.unfreeze(restored_train_state.params), sep="/"
      )
      del restored_train_state
      # Copy parameters over:
      for pname, pvalue in restored_params_dict.items():
          if "output_projection" in pname or "head" in pname:
              # Don't copy the output projection weight and bias from the checkpoint.
              continue
          elif (
              "MAPHead" in pname and model_cfg.model.classifier != "map"
          ):
              # Skip the MapHead parameters.
              continue
          elif "pos_embedding" in pname:
              params_dict[pname] = resample_posemb(
                  pvalue, params_dict[pname]
              )
          else:
              params_dict[pname] = pvalue
  
      logging.info(
          "Inspect missing keys from the restored params:\n%s",
          params_dict.keys() - restored_params_dict.keys(),
      )
      logging.info(
          "Inspect extra keys the the restored params:\n%s",
          restored_params_dict.keys() - params_dict.keys(),
      )
  
      # Restore data format, then initialize embeddings
      params = flax.traverse_util.unflatten_dict(params_dict, sep="/")
      return train_state.replace(params=flax.core.freeze(params))
  
imports: import functools
  from typing import Any, Optional
  from absl import logging
  import flax
  import flax.linen
  from flax.training import common_utils
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models import multilabel_classification_model
  from scenic.model_lib.layers import nn_layers
  import scipy
label_desc: Use image node_id 3 for calling a global function with example usage: init_vit_from_train_state(train_state, restored_train_state, model_cfg, restored_model_cfg) and returns: train_state