class_docstr: Vision Transformer model for multi-label classification task.
class_docstr_tok: ["Vision","Transformer","model","for","multi-label","classification","task","."]
docstr: Updates the train_state with data from restored_train_state.
  
  This function is writen 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:
    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.
  
  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",".","Here",",","we","do","some","surgery","to","support","larger","resolutions","(","longer","sequence","length",")","in","the","transformer","block",",","with","respect","to","the","learned","pos-embeddings",".","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",".","Returns",":","Updated","train_state","."]
code: def init_from_train_state(
      self,
      train_state: Any,
      restored_train_state: Any,
      restored_model_cfg: ml_collections.ConfigDict,
  ) -> Any:
      """Updates the train_state with data from restored_train_state.
  
      This function is writen 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:
        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.
  
      Returns:
        Updated train_state.
      """
      return vit.init_vit_from_train_state(
          train_state,
          restored_train_state,
          self.config,
          restored_model_cfg,
      )
  
imports: from typing import Any, Callable, Optional, Sequence
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import multilabel_classification_model
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.av_mae import base_model
  from scenic.projects.av_mae import model_utils
  from scenic.projects.baselines import vit
label_desc: Use image node_id 3 for calling the ViTMAEClassificationFinetuning obj's underlying member method code with example usage: obj.init_from_train_state(train_state, restored_train_state, restored_model_cfg) and returns: vit