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.
docstr: Add modality learned tokens.
docstr_tok: ["Add","modality","learned","tokens","."]
code: def add_modality_token(
      self, x_tokens_dict: ArrayDict, name: str = "Encoder"
  ) -> ArrayDict:
      """Add modality learned tokens."""
      if not self.use_modality_tokens:
          return x_tokens_dict
      for key, x_tokens in x_tokens_dict.items():
          modality_token = self.param(
              f"{name}_modality_token_{key}",
              nn.initializers.zeros,
              (1, 1, x_tokens.shape[-1]),
          )
          x_tokens = x_tokens + modality_token
          x_tokens_dict[key] = x_tokens
  
      return x_tokens_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 2 for calling the ViViTMultiMaskedAutoencoder obj's underlying member method code with example usage: obj.add_modality_token(x_tokens_dict, name) and returns: x_tokens_dict, x_tokens_dict