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: Shuffle the tokens and the token masks.
docstr_tok: ["Shuffle","the","tokens","and","the","token","masks","."]
code: def shuffle_tokens_and_token_mask(
      self,
      x_tokens_dict: ArrayDict,
      token_mask_dict: ArrayDict,
      rng: jax.Array,
  ) -> Tuple[ArrayDict, ArrayDict]:
      """Shuffle the tokens and the token masks."""
      # For the inpainting strategy, we add the unmasked token (from a different
      # modality) at the begining of the array and then we complete with unmasked
      # tokens. Now, we shuffle the tokens and the token masks.
      # TODO(lgeorgescu): hard assumption that RGB is in the dict.
      n_batch = x_tokens_dict[base_model.FeatureTargets.RGB].shape[0]
      batch_indices = jnp.arange(n_batch).reshape(n_batch, 1)
      rng_keys = jax.random.split(rng, n_batch * len(x_tokens_dict))
      idx_rng_key = 0
      for key in x_tokens_dict:
          x_tokens = x_tokens_dict[key]
          token_mask = token_mask_dict[key]
          n_tokens = x_tokens.shape[1]
          ids = jnp.tile(jnp.arange(n_tokens), n_batch).reshape(
              (n_batch, n_tokens)
          )
          ids = jax.vmap(
              lambda seq, rng: jax.random.permutation(
                  rng, seq, independent=True
              )
          )(
              ids,
              rng_keys[
                  idx_rng_key * n_batch : (idx_rng_key + 1) * n_batch
              ],
          )
  
          x_tokens = x_tokens.at[batch_indices, ids].set(x_tokens)
          token_mask = token_mask.at[batch_indices, ids].set(token_mask)
          x_tokens_dict[key] = x_tokens
          token_mask_dict[key] = token_mask
          idx_rng_key += 1
  
      return x_tokens_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 10 for calling the ViViTMultiMaskedAutoencoder obj's underlying member method code with example usage: obj.shuffle_tokens_and_token_mask(x_tokens_dict, token_mask_dict, rng) and returns: x_tokens_dict,  token_mask_dict