code: def initialize_from_vit_train_states(
      config: ml_collections.ConfigDict,
      train_state: train_utils.TrainState,
      restored_train_states: Sequence[train_utils.TrainState],
      restored_model_cfgs: Sequence[ml_collections.ConfigDict],
      restored_model_formats: Sequence[str],
  ) -> train_utils.TrainState:
      """Updates MTV's train_state with pretrained ViT weights.
  
      Args:
        config: Configurations for the model being updated.
        train_state: A raw TrainState for the model.
        restored_train_states: 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_cfgs: Configurations of models from which the
          restored_train_states come from. Often only the classifier information is
          used for interpolating the positional embeddings.
        restored_model_formats: A list of pretrained model formats. The format can
          only be `big_vision` or 'scenic'.
  
      Returns:
        Updated train_state.
      """
      assert len(restored_train_states) == len(
          restored_model_formats
      ), (
          "restored_train_states must have the same dimension as "
          "restored_model_formats."
      )
      params = flax.core.unfreeze(train_state.params)
      for view_idx, restored_state in enumerate(restored_train_states):
          restored_model_params = flax.core.unfreeze(
              restored_state.params
          )
          initialize_one_view_from_vit_parameters(
              config,
              params,
              restored_model_cfgs[view_idx],
              restored_model_params,
              view_idx,
          )
  
      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 11 for calling a global function with example usage: initialize_from_vit_train_states(config, train_state, restored_train_states, restored_model_cfgs, restored_model_formats) and returns: train_state