class_docstr: Vision Video Transformer model for MAE finetuning.
class_docstr_tok: ["Vision","Video","Transformer","model","for","MAE","finetuning","."]
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."""
      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."
          )
  
      # Move the encoder norm if it is there outside of the transformer encoder.
      if "encoder_norm" in restored_train_state.params["Transformer"]:
          restored_parameters = flax.core.unfreeze(
              restored_train_state.params
          )
          norm_parameters = restored_parameters["Transformer"].pop(
              "encoder_norm"
          )
          restored_parameters["encoder_norm"] = norm_parameters
          restored_train_state = restored_train_state.replace(
              params=flax.core.freeze(restored_parameters)
          )
  
      # If we restore from a non-MAE checkpoint and the positional embedding is
      # 'learned_1d', we have to move the positional embedding outside
      # the 'Transformer' block and also to drop the value of cls positional
      # embedding from the positonal embedding and add it to cls token.
  
      if (
          self.config.init_from.get(
              "restore_from_non_mae_checkpoint", False
          )
          and self.config.model.get(
              "positional_embedding", "sinusoidal_1d"
          )
          == "learned_1d"
      ):
          # Move the positional embedding outside the 'Transformer' block.
          restored_parameters = flax.core.unfreeze(
              restored_train_state.params
          )
          restored_parameters["posembed_input"] = restored_parameters[
              "Transformer"
          ].pop("posembed_input")
  
          if restored_model_cfg.model.classifier == "token":
              pos_embedding_params = restored_parameters[
                  "posembed_input"
              ]["pos_embedding"]
              # Drop the value of cls positional embedding.
              cls_pos_embedding = pos_embedding_params[:, 0]
              restored_parameters["posembed_input"][
                  "pos_embedding"
              ] = pos_embedding_params[:, 1:]
              # Add the value of cls positional embedding to cls token.
              restored_parameters["cls"] = (
                  restored_parameters["cls"] + cls_pos_embedding
              )
  
          restored_train_state = restored_train_state.replace(
              params=flax.core.freeze(restored_parameters)
          )
  
      return vivit_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: from typing import Any, Optional
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import classification_model
  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.vivit import model
  from scenic.projects.vivit import model_utils
label_desc: Use image node_id 2 for calling the ViViTMAEClassificationFinetuningModel 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