class_docstr: ViViT model for Multi-Modality Masked AutoEncoder.
  
  
  Attributes:
    mlp_dim: Dimension of the mlp on top of attention block.
    num_classes_dict: Dictionary with the number of output classes.
    num_heads: Number of self-attention heads.
    num_layers: Number of layers.
    patches: Configuration of the patches extracted in the stem of the model.
    hidden_size: Size of the hidden state of the output of model's stem.
    token_mask_probability_dict: Probability of dropping out the input tokens
    during training.
    masking_strategy: Masking strategy used to mask the tokens.
    temporal_encoding_config: ConfigDict which defines the type of input
      encoding when tokenising the video.
    decoder_config: ConfigDict which define the decoder.
    attention_config: ConfigDict which defines the type of spatio-temporal
      attention applied in the model.
    dropout_rate: Dropout rate.
    attention_dropout_rate: Dropout for attention heads.
    stochastic_droplayer_rate: Probability of dropping a layer. Linearly
      increases from 0 to the provided value..
    classifier: type of the classifier layer. Options are 'gap', 'gmp', 'gsp',
      'token'.
    encoder_strategy: Specify how to combine the modalities in the encoder.
      Choose from: separate_encoders, concat_and_encode, same_encoder.
    decoder_strategy: Specify how to combine the modalities in the decoder,
      Choose from: separate_decoders, same_decoder.
    use_inpainting: Whether or not to use the modality inpaiting strategy.
    normalise_encoder_output: If true, layer normalisation is applied to the
      output of the transformer encoder.
    use_modality_tokens: If True, modality learnable tokens are added.
    fusion_layers: When the encoder strategy is 'encode_and_concat',
      this specify how many layers to use for concatenation.
    dtype: JAX data type for activations.
code: def __call__(
      self, inputs: ArrayDict, *, train: bool, debug: bool = False
  ):
      del debug
      # inputs will be a dictionary
  
      x_tokens_dict = tokenize_input(
          inputs,
          temporal_encoding_config=self.temporal_encoding_config,
          patches=self.patches,
          hidden_size=self.hidden_size,
      )
  
      x_tokens_dict = add_positional_embeddings(
          x_tokens_dict, positional_embedding=self.positional_embedding
      )
  
      x_tokens_dict = self.add_modality_token(x_tokens_dict)
  
      (
          x_unmasked_dict,
          token_mask_dict,
          unmasked_indices_dict,
          masked_indices_dict,
      ) = self.mask_tokens(x_tokens_dict, train=train)
      # If we want to add a class token, add it here.
      # Note that in MAE, positional encodings are not added to the CLS token.
      x_unmasked_dict = add_cls_token(
          x_unmasked_dict, self.classifier, None
      )
  
      x_unmasked_encoded_dict = apply_encoder(
          x_unmasked_dict,
          encoder_strategy=self.encoder_strategy,
          mlp_dim=self.mlp_dim,
          num_layers=self.num_layers,
          num_heads=self.num_heads,
          attention_config=self.attention_config,
          dropout_rate=self.dropout_rate,
          attention_dropout_rate=self.attention_dropout_rate,
          stochastic_droplayer_rate=self.stochastic_droplayer_rate,
          dtype=self.dtype,
          normalise_encoder_output=self.normalise_encoder_output,
          fusion_layers=self.fusion_layers,
          train=train,
      )
  
      if not train:
          return x_unmasked_encoded_dict, token_mask_dict
  
      # Process entire sequence with the decoder.
      mask_token_dict = {}
      for key in x_unmasked_encoded_dict:
          mask_token_ = self.param(
              f"mask_token_{key}",
              nn.initializers.zeros,
              (1, 1, self.decoder_config.hidden_size),
          )
          mask_token_dict[key] = mask_token_
  
      x_unmasked_proj_dict = self.apply_decoder_projection(
          x_unmasked_encoded_dict
      )
  
      cls_encoded_dict, x_unmasked_proj_dict = self.remove_cls_token(
          x_unmasked_proj_dict
      )
  
      if self.use_inpainting:
          (
              x_all_dict,
              token_mask_dict,
          ) = self.unshuffle_tokens_inpainting(
              x_unmasked_proj_dict, x_tokens_dict, mask_token_dict
          )
      else:
          x_all_dict = self.unshuffle_tokens(
              x_unmasked_proj_dict,
              unmasked_indices_dict,
              masked_indices_dict,
              x_tokens_dict,
              mask_token_dict,
          )
  
      # Note. VideoMAE (Facebook) adds positional encodinggs to the CLS token at
      # the encoder as well. VideoMAE (Tong et al) don't use a CLS token.
      # This implementation does not add positional embeddings to the CLS token,
      # as in the original image MAE of He et al.
      x_all_dict = add_positional_embeddings(
          x_all_dict,
          positional_embedding=self.positional_embedding_decoder,
      )
      x_all_dict = self.add_modality_token(x_all_dict, name="Decoder")
  
      # Shuffle the tokens and token_mask accordingly if use_inpainting is True.
      if self.use_inpainting and self.shuffle_inpainted_tokens:
          (
              x_all_dict,
              token_mask_dict,
          ) = self.shuffle_tokens_and_token_mask(
              x_all_dict, token_mask_dict, rng=self.make_rng("dropout")
          )
  
      x_all_dict = add_cls_token(
          x_all_dict, self.classifier, cls_encoded_dict
      )
  
      x_decoded_dict = self.apply_decoder(x_all_dict, train=train)
  
      _, x_decoded_dict = self.remove_cls_token(x_decoded_dict)
  
      x_prelogits_dict = {}
      for key, x_decoded in x_decoded_dict.items():
          x_prelogits = nn_layers.IdentityLayer(
              name=f"pre_logits_{key}"
          )(x_decoded)
          x_prelogits_dict[key] = x_prelogits
  
      x_logits_dict = self.apply_dense_layer(x_prelogits_dict)
  
      return x_logits_dict, token_mask_dict
  
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 ViViTMultiMaskedAutoencoder obj's underlying member method code with example usage: obj.__call__(inputs) and returns: x_logits_dict,  token_mask_dict, x_unmasked_encoded_dict,  token_mask_dict