class_docstr: Video Transformer model for n-way classification.
class_docstr_tok: ["Video","Transformer","model","for","n-way","classification","."]
code: def build_flax_model(self) -> nn.Module:
      model_dtype = getattr(
          jnp, self.config.get("model_dtype_str", "float32")
      )
      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",
      ]:
          return ViViT(
              num_classes=self.dataset_meta_data["num_classes"],
              mlp_dim=self.config.model.mlp_dim,
              num_layers=self.config.model.num_layers,
              num_heads=self.config.model.num_heads,
              representation_size=self.config.model.representation_size,
              patches=self.config.model.patches,
              hidden_size=self.config.model.hidden_size,
              temporal_encoding_config=self.config.model.temporal_encoding_config,
              attention_config=self.config.model.attention_config,
              classifier=self.config.model.classifier,
              dropout_rate=self.config.model.get("dropout_rate", 0.1),
              attention_dropout_rate=self.config.model.get(
                  "attention_dropout_rate", 0.1
              ),
              stochastic_droplayer_rate=self.config.model.get(
                  "stochastic_droplayer_rate", 0
              ),
              return_prelogits=self.config.model.get(
                  "return_prelogits", False
              ),
              return_preclassifier=self.config.model.get(
                  "return_preclassifier", False
              ),
              dtype=model_dtype,
          )
      elif attention_type == "factorized_encoder":
          # TODO(dehghani): Rewrite this as a type of attention in ViViT Encoder.
          return SpaceTimeViViT(
              num_classes=self.dataset_meta_data["num_classes"],
              spatial_mlp_dim=self.config.model.spatial_transformer.mlp_dim,
              spatial_num_layers=self.config.model.spatial_transformer.num_layers,
              spatial_num_heads=self.config.model.spatial_transformer.num_heads,
              temporal_mlp_dim=self.config.model.temporal_transformer.mlp_dim,
              temporal_num_layers=self.config.model.temporal_transformer.num_layers,
              temporal_num_heads=self.config.model.temporal_transformer.num_heads,
              representation_size=self.config.model.representation_size,
              patches=self.config.model.patches,
              hidden_size=self.config.model.hidden_size,
              temporal_encoding_config=self.config.model.temporal_encoding_config,
              attention_config=self.config.model.attention_config,
              classifier=self.config.model.classifier,
              dropout_rate=self.config.model.get("dropout_rate", 0.1),
              attention_dropout_rate=self.config.model.get(
                  "attention_dropout_rate", 0.1
              ),
              stochastic_droplayer_rate=self.config.model.get(
                  "stochastic_droplayer_rate", 0
              ),
              return_prelogits=self.config.model.get(
                  "return_prelogits", False
              ),
              dtype=model_dtype,
          )
      else:
          raise ValueError(
              f"Attention type {attention_type} does not exist."
          )
  
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 1 for calling the ViViTClassificationModel obj's underlying member method code with example usage: obj.build_flax_model() and returns: ViViT, SpaceTimeViViT