code: def initialize_from_mtv_train_state(
      config: ml_collections.ConfigDict,
      train_state: train_utils.TrainState,
      restored_train_state: train_utils.TrainState,
      restored_model_cfg: ml_collections.ConfigDict,
      restore_output_projection: bool,
      model_prefix_path: Optional[List[str]] = None,
  ) -> train_utils.TrainState:
      """Updates MTV's train_state with a pretrained MTV weights.
  
      Args:
        config: Configurations for the model being updated.
        train_state: A raw TrainState for the model.
        restored_train_state: A dict. Each key is a Ti/S/B/L ViT model and the
          corresponding value is a TrainState that is loaded with parameters/state
          of pretrained models.
        restored_model_cfg: Configurations of models from which the
          restored_train_states come from. Often only the classifier information is
          used for interpolating the positional embeddings.
        restore_output_projection: Whether or not to restore the weights from output
          projection.
        model_prefix_path: The parent keys in the model dict where the restored
          model should reside.
  
      Returns:
        Updated train_state.
      """
  
      params = flax.core.unfreeze(train_state.params)
      restored_params = flax.core.unfreeze(restored_train_state.params)
      initialize_from_mtv_parameters(
          config,
          params,
          restored_model_cfg,
          restored_params,
          restore_output_projection,
          model_prefix_path,
      )
      return train_state.replace(params=flax.core.freeze(params))
  
imports: from typing import Any, Dict, List, Optional, Sequence
  from absl import logging
  import flax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.projects.vivit import model_utils
  from scenic.train_lib import train_utils
  import scipy.ndimage
label_desc: Use image node_id 9 for calling a global function with example usage: initialize_from_mtv_train_state(config, train_state, restored_train_state, restored_model_cfg, restore_output_projection, model_prefix_path) and returns: train_state