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:
      """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.
      """
      # Split up configs
      if one_config:
          init_config = config.init_from
          model_config = config.model
          dataset_config = config.dataset_configs
      else:
          init_config, model_config, dataset_config = config
  
      # Inspect and compare the parameters of the model with the init-model
      params = flax.core.unfreeze(train_state.optimizer.target)
      logging.info("Parameters in the target model are: %s", params)
  
      if init_config.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)
      if init_config.get("init_from_vit", True):
          if prefix_path:
              video_params = params[prefix_path]
          else:
              video_params = 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:
                      video_params[m_key] = m_params
                  else:
                      pass
              elif m_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.
                      video_params.pop(m_key, None)
                  else:
                      assert (
                          restored_model_cfg.model.representation_size
                      )
                      video_params[m_key] = m_params
  
              elif m_key in ["Transformer"]:
                  for tm_key, tm_params in m_params.items():
                      if (
                          tm_key == "posembed_input"
                      ):  # Might need resolution change
                          init_posemb(
                              video_params[mbt_transformer_key],
                              m_params,
                              init_config,
                              model_config,
                              dataset_config,
                              restored_model_cfg,
                              "posembed_input",
                              prefix_path=prefix_path,
                          )
                          init_posemb(
                              video_params[mbt_transformer_key],
                              m_params,
                              init_config,
                              model_config,
                              dataset_config,
                              restored_model_cfg,
                              "posembed_input_spectrogram",
                              prefix_path=prefix_path,
                          )
                          init_posemb(
                              video_params,
                              m_params,
                              init_config,
                              model_config,
                              dataset_config,
                              restored_model_cfg,
                              "bottleneck",
                              prefix_path=prefix_path,
                          )
                      elif "encoderblock" in tm_key:
                          logging.info("Loading encoder parameters.")
                          init_encoderblock(
                              video_params[mbt_transformer_key],
                              m_params,
                              tm_key,
                          )
                      else:  # Other parameters of the Transformer encoder
                          video_params[mbt_transformer_key][
                              tm_key
                          ] = tm_params
              elif m_key == "embedding":
                  init_embedding(
                      video_params,
                      m_params,
                      init_config,
                      model_config,
                      "embedding",
                  )
                  init_embedding(
                      video_params,
                      m_params,
                      init_config,
                      model_config,
                      "embedding_spectrogram",
                  )
              else:
                  if m_key in train_state.optimizer.target:
                      video_params[m_key] = m_params
                  if (
                      "%s_spectrogram" % m_key
                      in train_state.optimizer.target
                  ):
                      video_params["%s_spectrogram" % m_key] = m_params
                  else:
                      logging.info(
                          "Skipping %s. In restored model but not in target",
                          m_key,
                      )
      else:
          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 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(m_key, None)
                  else:
                      assert (
                          restored_model_cfg.model.representation_size
                      )
                      params[m_key] = m_params
              else:
                  if m_key in train_state.optimizer.target:
                      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(
          optimizer=train_state.optimizer.replace(
              target=flax.core.freeze(params)
          )
      )
  
file_path: scenic/scenic/projects/mbt/model_utils.py
imports: from typing import Any
  from absl import logging
  import flax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.common_lib import debug_utils
  from scenic.projects.vivit import model_utils
  import scipy
label_desc: Use image node_id 2 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