code: def initialise_from_train_state(
      config,
      train_state: Any,
      restored_train_state: Any,
      restored_model_cfg: ml_collections.ConfigDict,
      restore_output_proj: bool,
      mbt_transformer_key: str = "Transformer",
      log_initialised_param_shapes: bool = True,
      one_config: bool = True,
      prefix_path: Any = None,
  ) -> Any:
      # TODO(aarnab): Deal with Optax and flax.Optim format train states.
      """Updates the train_state with data from restored_train_state.
  
      This function is written to be used for 'fine-tuning' experiments. Here, we
      do some surgery to support larger resolutions (longer sequence length) in
      the transformer block, with respect to the learned pos-embeddings.
  
      Args:
        config: Configurations for the model being updated, or tuple of configs.
        train_state: A raw TrainState for the model.
        restored_train_state: A TrainState that is loaded with parameters/state of a
          pretrained model.
        restored_model_cfg: Configuration of the model from which the
          restored_train_state come from. Usually used for some asserts.
        restore_output_proj: If true, load the final output projection. Set
          to False if finetuning to a new dataset.
        mbt_transformer_key: The key used for storing the subtree in the
          parameters that keeps Transformer weights, that are supposed to be
          initialized from the given pre-trained model.
        log_initialised_param_shapes: If true, print tabular summary of all the
          variables in the model once they have been initialised.
        one_config: If true, we have only a single config. If false, we get a tuple
          of configs in the order [init_config, model_config, dataset_config]. This
          is useful for works that build upon MBT and have different models in their
          config.
        prefix_path: If parameters are in a subtree.
  
      Returns:
        Updated train_state.
      """
  
      def _get_optimizer(train_state):
          if hasattr(train_state, "optimizer"):
              if hasattr(train_state.optimizer, "target"):
                  return train_state.optimizer.target
              else:
                  return train_state.optimizer["target"]
          else:
              return train_state.params
  
      # Split up configs
      if one_config:
          init_config = config.init_from
          model_config = config.model
      else:
          init_config, model_config, _ = config
  
      params = flax.core.unfreeze(_get_optimizer(train_state))
      logging.info(
          "Parameters in the target model are: %s", params.keys()
      )
      restored_params = flax.core.unfreeze(
          _get_optimizer(restored_train_state)
      )
  
      if init_config.model_type == "vit":
          params = initialise_from_vit(
              params=params,
              restored_params=restored_params,
              config=config,
              restored_model_cfg=restored_model_cfg,
              mbt_transformer_key=mbt_transformer_key,
              restore_output_proj=restore_output_proj,
              prefix_path=prefix_path,
          )
  
      elif init_config.model_type == "multimae":
          encoder_strategy = restored_model_cfg.model.encoder_strategy
          if encoder_strategy == "separate_encoders":
              params = initialise_from_separate_encoders(
                  params, restored_params, config, restored_model_cfg
              )
          elif encoder_strategy == "separate_encoders_and_concat":
              params = initialise_from_mid_fusion(
                  params, restored_params, config, restored_model_cfg
              )
          elif encoder_strategy in [
              "concat_and_encode",
              "same_encoder",
          ]:
              params = initialise_from_same_encoder(
                  params, restored_params, config, restored_model_cfg
              )
          else:
              raise AssertionError(
                  f"Unsupported encoder strategy {encoder_strategy}."
              )
  
      elif init_config.model_type == "mbt":
          for m_key, m_params in restored_params.items():
              logging.info("mkey is: %s", m_key)
              if "ProjectionHeadAndClassifier" in m_key:
                  for tm_key, tm_params in m_params.items():
                      if tm_key == "output_projection":
                          if restore_output_proj:
                              params[m_key][tm_key] = tm_params
                          else:
                              logging.info(
                                  "Skipping output projection in restoring weights"
                              )
                              pass
                      elif tm_key == "pre_logits":
                          if model_config.representation_size is None:
                              # We don't have representation_size in the new model, so let's
                              # ignore if from the pretained model, in case it has it.
                              # Note, removing the key from the dictionary is necessary to
                              # prevent obscure errors from the Flax optimizer.
                              params.pop(tm_key, None)
                          else:
                              assert (
                                  restored_model_cfg.model.representation_size
                              )
                              params[m_key][tm_key] = tm_params
              else:
                  if m_key in params:
                      params[m_key] = m_params
                  else:
                      logging.info(
                          "Skipping %s. In restored model but not in target",
                          m_key,
                      )
      else:
          raise ValueError(
              f"Type of model initialising from unknown: {init_config.model_type}"
          )
  
      if log_initialised_param_shapes:
          logging.info(
              "Parameter summary after initialising from train state"
          )
          debug_utils.log_param_shapes(params)
      if hasattr(train_state, "optimizer"):
          return train_state.replace(
              optimizer=train_state.optimizer.replace(
                  target=flax.core.freeze(params)
              )
          )
      else:
          return train_state.replace(params=flax.core.freeze(params))
  
imports: import functools
  import re
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple
  from absl import logging
  import flax
  import flax.linen
  import jax
  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 classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.av_mae import model_utils
  from scenic.projects.baselines import vit
  from scenic.projects.mbt import model
  from scenic.projects.vivit import model_utils
  import scipy
label_desc: Use image node_id 3 for calling a global function with example usage: initialise_from_train_state(config, train_state, restored_train_state, restored_model_cfg, restore_output_proj, mbt_transformer_key, log_initialised_param_shapes, one_config, prefix_path) and returns: train_state, train_state, train_state, train_state, train_state