code: def initialise_from_vit_train_state(
      config,
      train_state: Any,
      restored_train_state: Any,
      restored_model_cfg: ml_collections.ConfigDict,
      log_initialised_param_shapes: bool = True,
  ) -> Any:
      """Updates the train_state with data from restored_train_state (ViT model).
  
      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.
        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.
      """
      # Inspect and compare the parameters of the model with the init-model.
      params = flax.core.unfreeze(train_state.params)
      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 in ["Transformer", "SpatialTransformer"]:
              for tm_key, tm_params in m_params.items():
                  if (
                      tm_key == "posembed_input"
                  ):  # Might need resolution change.
                      if "tokenizer2d" in params["tokenizer"]:
                          init_posemb(
                              params["tokenizer"]["tokenizer2d"],
                              m_params,
                              config,
                              restored_model_cfg,
                              "resize",
                          )
                      if "tokenizer3d" in params["tokenizer"]:
                          init_posemb(
                              params["tokenizer"]["tokenizer3d"],
                              m_params,
                              config,
                              restored_model_cfg,
                              config.init_from.positional_embed_size_change,
                          )
                      if "tokenizer_spec" in params["tokenizer"]:
                          init_spec_posemb(
                              params["tokenizer"]["tokenizer_spec"],
                              m_params,
                              config,
                              restored_model_cfg,
                          )
                  elif "encoderblock" in tm_key:
                      init_encoderblock(params, m_params, tm_key)
                  else:  # Other parameters of the Transformer encoder.
                      params["vit_encoder"][tm_key] = tm_params
          elif m_key == "cls":
              for tokenizer_name in [
                  "tokenizer2d",
                  "tokenizer3d",
                  "tokenizer_spec",
              ]:
                  if tokenizer_name in params["tokenizer"]:
                      params["tokenizer"][tokenizer_name][
                          "cls"
                      ] = m_params
          elif m_key == "embedding":
              for tokenizer_name in ["tokenizer2d", "tokenizer_spec"]:
                  if tokenizer_name in params["tokenizer"]:
                      params["tokenizer"][tokenizer_name][
                          "embedding"
                      ] = m_params
              if "tokenizer3d" in params["tokenizer"]:
                  init_embedding(
                      params["tokenizer"]["tokenizer3d"],
                      m_params,
                      config,
                  )
          else:
              if m_key in train_state.params:
                  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)
      return train_state.replace(params=flax.core.freeze(params))
  
file_path: scenic/scenic/projects/polyvit/model_utils.py
imports: from typing import Any, Optional
  from absl import logging
  import flax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.common_lib import debug_utils
  import scenic.projects.mbt.model_utils
  import scenic.projects.vivit.model_utils
  import scipy
label_desc: Use image node_id 4 for calling a global function with example usage: initialise_from_vit_train_state(config, train_state, restored_train_state, restored_model_cfg, log_initialised_param_shapes) and returns: train_state