class_docstr: MTV model for multiclass classification task.
docstr: Updates the train_state with data from restored_train_states.
  
  This function is used to initialize a MTV model from a list of ViT
  checkpoints. We assume that the number of restored_train_states is equal to
  the number of views.
  
  Args:
    train_state: A raw TrainState for the model.
    restored_train_states: A sequence of TrainStates that is loaded with
      parameters/state of a pretrained ViT model.
    restored_model_cfgs: A sequence of model configuration of the pretrained
      ViT models. Usually used for some asserts.
    restored_model_formats: The checkpoint format of each model. The format
      can be 'scenic' or 'big_vision'.
  
  Returns:
    Updated train_state.
docstr_tok: ["Updates","the","train_state","with","data","from","restored_train_states",".","This","function","is","used","to","initialize","a","MTV","model","from","a","list","of","ViT","checkpoints",".","We","assume","that","the","number","of","restored_train_states","is","equal","to","the","number","of","views",".","Args",":","train_state",":","A","raw","TrainState","for","the","model",".","restored_train_states",":","A","sequence","of","TrainStates","that","is","loaded","with","parameters\/state","of","a","pretrained","ViT","model",".","restored_model_cfgs",":","A","sequence","of","model","configuration","of","the","pretrained","ViT","models",".","Usually","used","for","some","asserts",".","restored_model_formats",":","The","checkpoint","format","of","each","model",".","The","format","can","be","'scenic","'","or","'big_vision","'",".","Returns",":","Updated","train_state","."]
code: def init_from_vit_train_states(
      self,
      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 the train_state with data from restored_train_states.
  
      This function is used to initialize a MTV model from a list of ViT
      checkpoints. We assume that the number of restored_train_states is equal to
      the number of views.
  
      Args:
        train_state: A raw TrainState for the model.
        restored_train_states: A sequence of TrainStates that is loaded with
          parameters/state of a pretrained ViT model.
        restored_model_cfgs: A sequence of model configuration of the pretrained
          ViT models. Usually used for some asserts.
        restored_model_formats: The checkpoint format of each model. The format
          can be 'scenic' or 'big_vision'.
  
      Returns:
        Updated train_state.
      """
      return model_utils.initialize_from_vit_train_states(
          self.config,
          train_state,
          restored_train_states,
          restored_model_cfgs,
          restored_model_formats,
      )
  
imports: import functools
  from typing import Any, List, Sequence, Tuple, Union
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
  from scenic.projects.mtv import model_utils
  from scenic.projects.vivit import model
  from scenic.train_lib import train_utils
label_desc: Use image node_id 4 for calling the MTVClassificationModel obj's underlying member method code with example usage: obj.init_from_vit_train_states(train_state, restored_train_states, restored_model_cfgs, restored_model_formats) and returns: model_utils