class_docstr: MTV model for multiclass classification task.
docstr: Updates the train_state with data from restored_train_state.
  
  This function is writen to be used for 'fine-tuning' experiments. The input
  embeddings and positional embeddings are resized if the current model uses
  a different size of tubelets than the pretrained model.
  
  Args:
    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: Whether or not to restore output projection weights.
  
  Returns:
    Updated train_state.
docstr_tok: ["Updates","the","train_state","with","data","from","restored_train_state",".","This","function","is","writen","to","be","used","for","'fine-tuning","'","experiments",".","The","input","embeddings","and","positional","embeddings","are","resized","if","the","current","model","uses","a","different","size","of","tubelets","than","the","pretrained","model",".","Args",":","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",":","Whether","or","not","to","restore","output","projection","weights",".","Returns",":","Updated","train_state","."]
code: def init_from_train_state(
      self,
      train_state: train_utils.TrainState,
      restored_train_state: train_utils.TrainState,
      restored_model_cfg: ml_collections.ConfigDict,
      restore_output_proj: bool = False,
  ) -> train_utils.TrainState:
      """Updates the train_state with data from restored_train_state.
  
      This function is writen to be used for 'fine-tuning' experiments. The input
      embeddings and positional embeddings are resized if the current model uses
      a different size of tubelets than the pretrained model.
  
      Args:
        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: Whether or not to restore output projection weights.
  
      Returns:
        Updated train_state.
      """
      return model_utils.initialize_from_mtv_train_state(
          self.config,
          train_state,
          restored_train_state,
          restored_model_cfg,
          restore_output_projection=restore_output_proj,
      )
  
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 3 for calling the MTVClassificationModel obj's underlying member method code with example usage: obj.init_from_train_state(train_state, restored_train_state, restored_model_cfg, restore_output_proj) and returns: model_utils