code: def initialise_from_train_state(
      train_state: Any,
      restored_train_state: Any,
      config: ml_collections.ConfigDict,
      restored_model_cfg: ml_collections.ConfigDict,
  ) -> Any:
      """Updates the train_state with data from restored_train_state.
  
      Args:
        train_state: A raw TrainState for the model.
        restored_train_state: A TrainState that is loaded with parameters/state of a
          pretrained model.
        config: Configurations for the model being updated.
        restored_model_cfg: Configuration of the model from which the
          restored_train_state come from. Usually used for some asserts.
  
      Returns:
        Updated train_state.
      """
      del config
      del restored_model_cfg
      # Create restored parameters dict:
      params = flax.core.unfreeze(train_state.optimizer.target)
      params_dict = {
          "/".join([str(kk) for kk in k]): v
          for k, v in flax.traverse_util.flatten_dict(params).items()
      }
      # Create restored parameters dict:
      restored_params = flax.core.unfreeze(
          restored_train_state.optimizer.target
      )
      restored_params_dict = dict()
      for key, value in flax.traverse_util.flatten_dict(
          restored_params
      ).items():
          name = "/".join([str(k) for k in key])
          restored_params_dict[name] = value
      # Copy parameters over:
      for pname, pvalue in restored_params_dict.items():
          if "output_projection" in pname:
              continue
          elif "pos_embedding" in pname:
              # TODO(dehghani) add support for reshaping pos-embedding to longer seq
              #  (e.g., for high res finetuning.).
              continue
          elif pname in params_dict:
              params_dict[pname] = pvalue
          else:
              logging.error(
                  "Restored key doesn't exist in the model: %s.", pname
              )
      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
      splitkeys = {
          tuple(k.split("/")): v for k, v in params_dict.items()
      }
      params = flax.traverse_util.unflatten_dict(splitkeys)
      logging.info(
          "Parameter summary after initialising from train state:"
      )
      debug_utils.log_param_shapes(params)
      return train_state.replace(
          optimizer=train_state.optimizer.replace(
              target=flax.core.freeze(params)
          )
      )
  
imports: 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.common_lib import debug_utils
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.adatape import layers
  from scenic.projects.baselines import vit
label_desc: Use image node_id 2 for calling a global function with example usage: initialise_from_train_state(train_state, restored_train_state, config, restored_model_cfg) and returns: train_state