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: Apply the decoder for each modality.
docstr_tok: ["Apply","the","decoder","for","each","modality","."]
code: def apply_decoder(
      self, x_all_dict: ArrayDict, train: bool
  ) -> ArrayDict:
      """Apply the decoder for each modality."""
  
      def concat_and_decode():
          x_decoded_dict = {}
          decoder = vivit_model.Encoder(
              temporal_dims=None,
              mlp_dim=self.decoder_config.mlp_dim,
              num_layers=self.decoder_config.num_layers,
              num_heads=self.decoder_config.num_heads,
              attention_config=self.decoder_config.attention_config,
              dropout_rate=self.decoder_config.dropout_rate,
              attention_dropout_rate=self.decoder_config.attention_dropout_rate,
              stochastic_droplayer_rate=(
                  self.decoder_config.stochastic_droplayer_rate
              ),
              dtype=self.dtype,
              positional_embedding="none",  # Has already been added.
              normalise_output=self.normalise_encoder_output,
              name="Decoder",
          )
  
          x_all_input_list = []
          x_all_key_list = []
          for key, x_all in x_all_dict.items():
              x_all_input_list.append(x_all)
              x_all_key_list.append(key)
  
          x_all_concat_input = jnp.concatenate(x_all_input_list, axis=1)
          x_all_decoded_concat = decoder(
              x_all_concat_input, train=train
          )
          start_index = 0
          for key, x_all_ in zip(x_all_key_list, x_all_input_list):
              end_index = start_index + x_all_.shape[1]
              x_decoded_dict[key] = x_all_decoded_concat[
                  :, start_index:end_index
              ]
              start_index = end_index
              assert x_decoded_dict[key].shape[1] == x_all_.shape[1]
  
          return x_decoded_dict
  
      def same_decoder():
          x_decoded_dict = {}
          decoder = vivit_model.Encoder(
              temporal_dims=None,
              mlp_dim=self.decoder_config.mlp_dim,
              num_layers=self.decoder_config.num_layers,
              num_heads=self.decoder_config.num_heads,
              attention_config=self.decoder_config.attention_config,
              dropout_rate=self.decoder_config.dropout_rate,
              attention_dropout_rate=self.decoder_config.attention_dropout_rate,
              stochastic_droplayer_rate=(
                  self.decoder_config.stochastic_droplayer_rate
              ),
              dtype=self.dtype,
              positional_embedding="none",  # Has already been added.
              normalise_output=self.normalise_encoder_output,
              name="Decoder",
          )
  
          for key, x_all in x_all_dict.items():
              x_decoded = decoder(x_all, train=train)
              x_decoded_dict[key] = x_decoded
          return x_decoded_dict
  
      def separate_decoders():
          x_decoded_dict = {}
          for key, x_all in x_all_dict.items():
              x_decoded = vivit_model.Encoder(
                  temporal_dims=None,
                  mlp_dim=self.decoder_config.mlp_dim,
                  num_layers=self.decoder_config.num_layers,
                  num_heads=self.decoder_config.num_heads,
                  attention_config=self.decoder_config.attention_config,
                  dropout_rate=self.decoder_config.dropout_rate,
                  attention_dropout_rate=self.decoder_config.attention_dropout_rate,
                  stochastic_droplayer_rate=(
                      self.decoder_config.stochastic_droplayer_rate
                  ),
                  dtype=self.dtype,
                  positional_embedding="none",  # Has already been added.
                  normalise_output=self.normalise_encoder_output,
                  name=f"Decoder_{key}",
              )(x_all, train=train)
              x_decoded_dict[key] = x_decoded
          return x_decoded_dict
  
      if self.decoder_strategy == "separate_decoders":
          return separate_decoders()
      elif self.decoder_strategy == "same_decoder":
          return same_decoder()
      elif self.decoder_strategy == "concat_and_decode":
          return concat_and_decode()
      else:
          raise ValueError(
              f"The decoder strategy {self.decoder_strategy} is not supported!"
          )
  
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 4 for calling the ViViTMultiMaskedAutoencoder obj's underlying member method code with example usage: obj.apply_decoder(x_all_dict, train) and returns: x_decoded_dict, x_decoded_dict, x_decoded_dict, separate_decoders, same_decoder, concat_and_decode