code: def initialise_from_train_state(
      config,
      train_state: Any,
      restored_train_state: Any,
      restored_model_cfg: ml_collections.ConfigDict,
      restore_output_proj: bool,
      log_initialised_param_shapes: bool = True,
      one_config: 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, 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.
        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.
  
      Returns:
        Updated train_state.
      """
      # Split up configs
      if one_config:
          init_config = config.init_from
          model_config = config.model
      else:
          init_config, model_config = config
  
      # Inspect and compare the parameters of the model with the init-model
      params = flax.core.unfreeze(train_state.optimizer.target)
  
      if init_config.get("checkpoint_format", "scenic") == "bigvision":
          restored_params = restored_train_state.optimizer["target"]
      else:
          restored_params = restored_train_state.optimizer.target
      restored_params = flax.core.unfreeze(restored_params)
      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.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/vid2seq/load_utils.py
imports: import copy
  import os
  import re
  from typing import List, Mapping, Optional, Union, Any
  from absl import logging
  import flax
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import ml_collections
  from scenic.common_lib import debug_utils
  from scenic.train_lib_deprecated.pretrain_utils import get_params_and_model_state_dict
  from scenic.train_lib_deprecated.pretrain_utils import inspect_params
  from scenic.train_lib_deprecated.train_utils import TrainState
  from tensorflow.io import gfile
label_desc: Use image node_id 4 for calling a global function with example usage: initialise_from_train_state(config, train_state, restored_train_state, restored_model_cfg, restore_output_proj, log_initialised_param_shapes, one_config) and returns: train_state