class_docstr: ViViT model for multi-modalitymasked autoencoder pretraining.
docstr: Updates the train_state with data from restored_train_state.
docstr_tok: ["Updates","the","train_state","with","data","from","restored_train_state","."]
code: def init_from_train_state(
      self,
      train_state: Any,
      restored_train_state: Any,
      restored_model_cfg: ml_collections.ConfigDict,
      restore_output_proj: bool = False,
  ) -> Any:
      """Updates the train_state with data from restored_train_state."""
      return vivit_model_utils.initialise_from_train_state(
          self.config,
          train_state,
          restored_train_state,
          restored_model_cfg,
          restore_output_proj,
      )
  
imports: import functools
  from typing import Any, Dict, Optional, Tuple
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  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.mbt import model
  from scenic.projects.vivit import model
  from scenic.projects.vivit import model_utils
label_desc: Use image node_id 2 for calling the ViViTMultiMaskedAutoencoderModel 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: vivit_model_utils