class_docstr: ViViT model for multi-modalitymasked autoencoder pretraining.
code: def build_flax_model(self) -> nn.Module:
      model_dtype = getattr(
          jnp, self.config.get("model_dtype_str", "float32")
      )
  
      num_classes_dict = {}
      for feature_target in self.config.masked_feature_loss.target:
          if feature_target == base_model.FeatureTargets.RGB:
              select_central_frame = (
                  self.config.masked_feature_loss.select_central_frame
              )
              patch_size = tuple(self.config.model.patches.size)
              channels = 3
          elif feature_target == base_model.FeatureTargets.SPECTROGRAM:
              patch_size = tuple(self.config.model.patches.size[:2])
              select_central_frame = False
              channels = 1
          else:
              raise ValueError(f"{feature_target} is not supported!")
  
          num_classes = base_model.get_output_shapes(
              feature_target,
              patch_size,
              select_central_frame,
              channels=channels,
          )
          num_classes_dict[feature_target] = num_classes
  
      return ViViTMultiMaskedAutoencoder(
          num_classes_dict=num_classes_dict,
          mlp_dim=self.config.model.mlp_dim,
          num_layers=self.config.model.num_layers,
          num_heads=self.config.model.num_heads,
          patches=self.config.model.patches,
          hidden_size=self.config.model.hidden_size,
          token_mask_probability_dict=(
              self.config.masked_feature_loss.token_mask_probability_dict
          ),
          masking_strategy=self.config.masked_feature_loss.get(
              "masking_strategy", "random"
          ),
          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.0),
          attention_dropout_rate=self.config.model.get(
              "attention_dropout_rate", 0.1
          ),
          stochastic_droplayer_rate=self.config.model.get(
              "stochastic_droplayer_rate", 0
          ),
          dtype=model_dtype,
          decoder_config=self.config.model.get("decoder_config", None),
          positional_embedding=self.config.model.get(
              "positional_embedding", "sinusoidal_1d"
          ),
          positional_embedding_decoder=self.config.model.get(
              "positional_embedding_decoder", "sinusoidal_1d"
          ),
          normalise_encoder_output=self.config.model.get(
              "normalise_encoder_output", True
          ),
          use_inpainting=self.config.model.get("use_inpainting", False),
          shuffle_inpainted_tokens=self.config.model.get(
              "shuffle_inpainted_tokens", False
          ),
          encoder_strategy=self.config.model.get("encoder_strategy"),
          decoder_strategy=self.config.model.get("decoder_strategy"),
          use_modality_tokens=self.config.model.get(
              "use_modality_tokens", False
          ),
          fusion_layers=self.config.model.get("fusion_layers", None),
      )
  
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 1 for calling the ViViTMultiMaskedAutoencoderModel obj's underlying member method code with example usage: obj.build_flax_model() and returns: ViViTMultiMaskedAutoencoder