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: Mask the tokens based on their probability.
docstr_tok: ["Mask","the","tokens","based","on","their","probability","."]
code: def mask_tokens(
      self, x_tokens_dict: ArrayDict, train: bool
  ) -> Tuple[ArrayDict, ArrayDict, ArrayDict, ArrayDict]:
      """Mask the tokens based on their probability."""
      x_unmasked_dict = {}
      token_mask_dict = {}
      unmasked_indices_dict = {}
      masked_indices_dict = {}
  
      for key, x_tokens in x_tokens_dict.items():
          n_batch, n_tokens = x_tokens.shape[:2]
          if train:
              if self.masking_strategy == "random":
                  # Generate mask indices by randomly masking the tokens.
                  n_masked = int(
                      self.token_mask_probability_dict[key] * n_tokens
                  )
                  (
                      mask_indices,
                      unmasked_indices,
                      token_mask,
                  ) = model_utils.get_mask_indices(
                      n_batch,
                      n_tokens,
                      n_masked,
                      self.make_rng("dropout"),
                  )
              else:
                  raise ValueError(
                      f"The masking strategy {self.masking_strategy} is not supported."
                  )
              # Process only unmasked tokens with the encoder.
              batch_indices = jnp.arange(n_batch).reshape(n_batch, 1)
              x_unmasked = x_tokens[batch_indices, unmasked_indices]
          else:
              x_unmasked = x_tokens
              token_mask = jnp.zeros((n_batch, n_tokens))
              # We won't need this if train is False.
              unmasked_indices = None
              mask_indices = None
  
          x_unmasked_dict[key] = x_unmasked
          token_mask_dict[key] = token_mask
          unmasked_indices_dict[key] = unmasked_indices
          masked_indices_dict[key] = mask_indices
  
      return (
          x_unmasked_dict,
          token_mask_dict,  # pytype: disable=bad-return-type  # jax-ndarray
          unmasked_indices_dict,
          masked_indices_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 9 for calling the ViViTMultiMaskedAutoencoder obj's underlying member method code with example usage: obj.mask_tokens(x_tokens_dict, train) and returns: x_unmasked_dict,  token_mask_dict,  unmasked_indices_dict,  masked_indices_dict