code: def initialise_from_train_state(
      config,
      train_state: Any,
      restored_train_state: Any,
      restored_model_cfg: ml_collections.ConfigDict,
      restore_output_proj: bool,
      vivit_transformer_key: str = "Transformer",
      log_initialised_param_shapes: bool = True,
  ) -> Any:
      """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.
        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.
        vivit_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.
  
      Returns:
        Updated train_state.
      """
      if hasattr(train_state, "optimizer"):
          # Inspect and compare the parameters of the model with the init-model.
          params = flax.core.unfreeze(train_state.optimizer.target)
          train_state_keys = train_state.optimizer.target.keys()
      else:
          params = flax.core.unfreeze(train_state.params)
          train_state_keys = train_state.params.keys()
      if hasattr(restored_train_state, "optimizer"):
          if (
              config.init_from.get("checkpoint_format", "scenic")
              == "big_vision"
          ):
              restored_params = restored_train_state.optimizer["target"]
          else:
              restored_params = restored_train_state.optimizer.target
          restored_params = flax.core.unfreeze(restored_params)
      else:
          restored_params = flax.core.unfreeze(
              restored_train_state.params
          )
  
      # Start moving parameters, one-by-one and apply changes if needed.
      for m_key, m_params in restored_params.items():
          if m_key == "output_projection":
              if restore_output_proj:
                  params[m_key] = m_params
              else:
                  pass
  
          elif m_key == "pre_logits":
              if config.model.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(m_key, None)
              else:
                  assert restored_model_cfg.model.representation_size
                  params[m_key] = m_params
  
          elif m_key in {
              "Transformer",
              "SpatialTransformer",
              "TemporalTransformer",
          }:
              key_to_load = vivit_transformer_key
              is_temporal = False
              if m_key == "TemporalTransformer":
                  key_to_load = m_key
                  is_temporal = True
              for tm_key, tm_params in m_params.items():
                  if (
                      tm_key == "posembed_input"
                  ):  # Might need resolution change.
                      init_posemb(
                          params[key_to_load],
                          m_params,
                          config,
                          restored_model_cfg,
                          is_temporal=is_temporal,
                      )
                  elif "encoderblock" in tm_key:
                      init_encoderblock(
                          params[key_to_load], m_params, tm_key, config
                      )
                  else:  # Other parameters of the Transformer encoder.
                      params[key_to_load][tm_key] = tm_params
  
          elif m_key == "embedding":
              init_embedding(params, m_params, config)
          else:
              if m_key in train_state_keys:
                  params[m_key] = m_params
              else:
                  logging.info(
                      "Skipping %s. In restored model but not in target",
                      m_key,
                  )
  
      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: from typing import Any, Optional, Tuple
  from absl import logging
  import flax
  from flax.linen import linear
  from flax.training import common_utils
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.common_lib import debug_utils
  from scenic.model_lib.base_models import model_utils
  import scipy
label_desc: Use image node_id 9 for calling a global function with example usage: initialise_from_train_state(config, train_state, restored_train_state, restored_model_cfg, restore_output_proj, vivit_transformer_key, log_initialised_param_shapes) and returns: train_state, train_state