class_docstr: Video Transformer model for n-way classification.
class_docstr_tok: ["Video","Transformer","model","for","n-way","classification","."]
docstr: 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."""
      attention_type = self.config.model.attention_config.get(
          "type", "spacetime"
      )
      if attention_type in [
          "spacetime",
          "factorized_transformer_block",
          "factorized_self_attention_block",
          "factorized_dot_product_attention",
      ]:
          vivit_transformer_key = "Transformer"
      elif attention_type == "factorized_encoder":
          vivit_transformer_key = "SpatialTransformer"
      else:
          raise ValueError(
              f"Attention type {attention_type} does not exist."
          )
      return model_utils.initialise_from_train_state(
          self.config,
          train_state,
          restored_train_state,
          restored_model_cfg,
          restore_output_proj,
          vivit_transformer_key=vivit_transformer_key,
      )
  
imports: import functools
  from typing import Any, Optional, Callable, Sequence
  from absl import logging
  import flax.linen
  from flax.linen.linear import default_kernel_init
  from immutabledict import immutabledict
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.common_lib import video_utils
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  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.vivit import model_utils
label_desc: Use image node_id 3 for calling the ViViTClassificationModel 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