class_docstr: X-ViT model for classification task.
docstr: Updates the train_state with data from restored_train_state.
  
  This function is written 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 written 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,
      )
  
file_path: scenic/scenic/projects/tasseo/xvit.py
imports: from typing import Any
  import ml_collections
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.projects.baselines import vit
  from scenic.projects.fast_vit.xvit import XViT
label_desc: Use image node_id 3 for calling the XViTClassificationModel 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