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: "Unshuffle the tokens for modality inpainting.
  
  Place the masked tokens of one modality (target) and the unmasked tokens
  of another modality in a matrix to create the input for the decoder.
  The unmasked tokens are placed first followed by the masked tokens.
  
  Args:
    x_unmasked_proj_dict: Dictionary with the unmasked tokens. The shape of
      the unmasked tokens is: [n_batch, n_unmasked_tokens, hidden_size].
    x_tokens_dict: Dictionary with all tokens. Used only for computing the
      total number of tokens. The shape of the tokens is:
      [n_batch, n_tokens, hidden_size].
    mask_token_dict: Dictionary with the masked tokens. The shape of the mask
      token is: [1, 1, hidden_size].
  
  Returns:
    x_all_dict: Dictionary with the input for the decoder. The shape of the
    tokens is: [n_batch, n_tokens, hidden_size].
docstr_tok: ["``","Unshuffle","the","tokens","for","modality","inpainting",".","Place","the","masked","tokens","of","one","modality","(","target",")","and","the","unmasked","tokens","of","another","modality","in","a","matrix","to","create","the","input","for","the","decoder",".","The","unmasked","tokens","are","placed","first","followed","by","the","masked","tokens",".","Args",":","x_unmasked_proj_dict",":","Dictionary","with","the","unmasked","tokens",".","The","shape","of","the","unmasked","tokens","is",":","[","n_batch",",","n_unmasked_tokens",",","hidden_size","]",".","x_tokens_dict",":","Dictionary","with","all","tokens",".","Used","only","for","computing","the","total","number","of","tokens",".","The","shape","of","the","tokens","is",":","[","n_batch",",","n_tokens",",","hidden_size","]",".","mask_token_dict",":","Dictionary","with","the","masked","tokens",".","The","shape","of","the","mask","token","is",":","[","1",",","1",",","hidden_size","]",".","Returns",":","x_all_dict",":","Dictionary","with","the","input","for","the","decoder",".","The","shape","of","the","tokens","is",":","[","n_batch",",","n_tokens",",","hidden_size","]","."]
code: def unshuffle_tokens_inpainting(
      self,
      x_unmasked_proj_dict: ArrayDict,
      x_tokens_dict: ArrayDict,
      mask_token_dict: ArrayDict,
  ) -> Tuple[ArrayDict, ArrayDict]:
      """ "Unshuffle the tokens for modality inpainting.
  
      Place the masked tokens of one modality (target) and the unmasked tokens
      of another modality in a matrix to create the input for the decoder.
      The unmasked tokens are placed first followed by the masked tokens.
  
      Args:
        x_unmasked_proj_dict: Dictionary with the unmasked tokens. The shape of
          the unmasked tokens is: [n_batch, n_unmasked_tokens, hidden_size].
        x_tokens_dict: Dictionary with all tokens. Used only for computing the
          total number of tokens. The shape of the tokens is:
          [n_batch, n_tokens, hidden_size].
        mask_token_dict: Dictionary with the masked tokens. The shape of the mask
          token is: [1, 1, hidden_size].
  
      Returns:
        x_all_dict: Dictionary with the input for the decoder. The shape of the
        tokens is: [n_batch, n_tokens, hidden_size].
      """
  
      def get_different_key(key_):
          available_keys = set(
              [
                  base_model.FeatureTargets.RGB,
                  base_model.FeatureTargets.SPECTROGRAM,
              ]
          )
          return next(iter((available_keys - set([key_]))))
  
      x_all_dict = {}
      token_mask_dict = {}
      for key in x_unmasked_proj_dict.keys():
          x_unmasked_proj = x_unmasked_proj_dict[get_different_key(key)]
          n_batch = x_unmasked_proj.shape[0]
          batch_indices = jnp.arange(n_batch).reshape(n_batch, 1)
          n_tokens = x_tokens_dict[key].shape[1]
          x_all = jnp.zeros(
              (n_batch, n_tokens, self.decoder_config.hidden_size)
          )
  
          # These are indices that come from the other modality. Note that as shapes
          # between modalities can vary, there is no correspondence between the
          # tokens. Therefore, we place all the unmasked tokens first in the token
          # sequence.
          num_unmasked_tokens = x_unmasked_proj.shape[1]
          unmasked_indices = jnp.repeat(
              jnp.arange(num_unmasked_tokens).reshape(
                  1, num_unmasked_tokens
              ),
              n_batch,
              axis=0,
          )
          x_all = x_all.at[batch_indices, unmasked_indices].set(
              x_unmasked_proj
          )
  
          # These are indices that will be updated with the current mask token.
          num_masked_tokens = n_tokens - num_unmasked_tokens
          masked_indices = jnp.repeat(
              jnp.arange(num_unmasked_tokens, n_tokens).reshape(
                  1, num_masked_tokens
              ),
              n_batch,
              axis=0,
          )
          x_all = x_all.at[batch_indices, masked_indices].set(
              mask_token_dict[key]
          )
  
          unmasked_token_mask = jnp.zeros(
              (n_batch, num_unmasked_tokens)
          )
          masked_token_mask = jnp.ones((n_batch, num_masked_tokens))
          token_mask = jnp.concatenate(
              (unmasked_token_mask, masked_token_mask), axis=1
          )
  
          x_all_dict[key] = x_all
          token_mask_dict[key] = token_mask
  
      return x_all_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 5 for calling the ViViTMultiMaskedAutoencoder obj's underlying member method code with example usage: obj.unshuffle_tokens_inpainting(x_unmasked_proj_dict, x_tokens_dict, mask_token_dict) and returns: x_all_dict,  token_mask_dict, next