class_docstr: Vision Video Transformer model for masked-autoencoding.
  
    The differences to the scenic.project.vivit.model.ViViT are that:
    -- Remove masked tokens from the encoder.
    -- Process all tokens with the decoder.
    -- The CLS token is added after the positional embeddings. This follows MAE.
  
  Attributes:
    mlp_dim: Dimension of the mlp on top of attention block.
    num_classes: 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: 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.
    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'.
    normalise_encoder_output: If true, layer normalisation is applied to the
      output of the transformer encoder.
    dtype: JAX data type for activations.
class_docstr_tok: ["Vision","Video","Transformer","model","for","masked-autoencoding",".","The","differences","to","the","scenic.project.vivit.model.ViViT","are","that",":","--","Remove","masked","tokens","from","the","encoder",".","--","Process","all","tokens","with","the","decoder",".","--","The","CLS","token","is","added","after","the","positional","embeddings",".","This","follows","MAE",".","Attributes",":","mlp_dim",":","Dimension","of","the","mlp","on","top","of","attention","block",".","num_classes",":","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",":","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",".","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","'",".","normalise_encoder_output",":","If","true",",","layer","normalisation","is","applied","to","the","output","of","the","transformer","encoder",".","dtype",":","JAX","data","type","for","activations","."]
code: def __call__(
      self, inputs: jnp.ndarray, *, train: bool, debug: bool = False
  ):
      del debug
      # Shape is [batch, num_tokens, hidden_size].
      x_tokens, temporal_dims = vivit_model.temporal_encode(
          inputs,
          self.temporal_encoding_config,
          self.patches,
          self.hidden_size,
      )
  
      n_batch, n_tokens, hidden_dim = x_tokens.shape
      height = width = int(np.sqrt(n_tokens // temporal_dims))
      if height * width * temporal_dims != n_tokens:
          raise ValueError("Input is assumed to be square.")
  
      # Add positional encodings.
      input_shape = None
      if self.positional_embedding == "sinusoidal_3d":
          input_shape = [
              n_batch,
              temporal_dims,
              height,
              width,
              hidden_dim,
          ]
      elif self.positional_embedding == "learned_space_time":
          input_shape = [
              n_batch,
              temporal_dims,
              height * width,
              hidden_dim,
          ]
  
      x_tokens = model_utils.add_positional_embeddings(
          x_tokens, self.positional_embedding, input_shape=input_shape
      )
  
      if train:
          if self.masking_strategy == "random":
              # Generate mask indices by randomly masking the tokens.
              n_masked = int(self.token_mask_probability * n_tokens)
              (
                  mask_indices,
                  unmasked_indices,
                  token_mask,
              ) = model_utils.get_mask_indices(
                  n_batch, n_tokens, n_masked, self.make_rng("dropout")
              )
  
          elif self.masking_strategy == "tube":
              # Generate mask indices by using tube masking.
              (
                  mask_indices,
                  unmasked_indices,
                  token_mask,
              ) = model_utils.get_tube_mask_indices(
                  n_batch=n_batch,
                  n_tokens=n_tokens,
                  token_mask_probability=self.token_mask_probability,
                  temporal_dims=temporal_dims,
                  rng=self.make_rng("dropout"),
              )
          else:
              raise ValueError(
                  f"The masking strategy {self.masking_strategy} is not implemented."
              )
          # 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))
  
      # If we want to add a class token, add it here.
      # Note that in MAE, positional encodings are not added to the CLS token.
      if self.classifier == "token":
          cls = self.param(
              "cls",
              nn.initializers.zeros,
              (1, 1, hidden_dim),
              inputs.dtype,
          )
          cls = jnp.tile(cls, [n_batch, 1, 1])
          x_unmasked = jnp.concatenate([cls, x_unmasked], axis=1)
  
      x_unmasked_encoded = vivit_model.Encoder(
          temporal_dims=temporal_dims,
          mlp_dim=self.mlp_dim,
          num_layers=self.num_layers,
          num_heads=self.num_heads,
          attention_config=self.attention_config,
          dropout_rate=self.dropout_rate,
          attention_dropout_rate=self.attention_dropout_rate,
          stochastic_droplayer_rate=self.stochastic_droplayer_rate,
          dtype=self.dtype,
          positional_embedding="none",  # Has already been added.
          normalise_output=self.normalise_encoder_output,
          name="Transformer",
      )(x_unmasked, train=train)
  
      x_unmasked_encoded = nn_layers.IdentityLayer(
          name="encoder_output"
      )(x_unmasked_encoded)
      if not train:
          x_representation = nn_layers.IdentityLayer(
              name="representation"
          )(x_unmasked_encoded)
          return x_representation, token_mask
  
      # Process entire sequence with the decoder.
      mask_token = self.param(
          "mask_token",
          nn.initializers.zeros,
          (1, 1, self.decoder_config.hidden_size),
      )
      x_unmasked_proj = nn.Dense(
          self.decoder_config.hidden_size,
          use_bias=self.decoder_config.get("use_projection_bias", True),
          kernel_init=nn.initializers.xavier_uniform(),
          name="decoder_projection",
      )(x_unmasked_encoded)
      if self.classifier == "token":
          cls_encoded = x_unmasked_proj[:, :1, :]
          x_unmasked_proj = x_unmasked_proj[:, 1:, :]
  
      # This effectively "unshuffles" the tokens. This means that we can simply
      # add positional encodings in the decoder without having to worry about
      # their ordering.
      x_all = jnp.zeros(
          (n_batch, n_tokens, self.decoder_config.hidden_size)
      )
      x_all = x_all.at[batch_indices, unmasked_indices].set(
          x_unmasked_proj
      )
      x_all = x_all.at[batch_indices, mask_indices].set(mask_token)
  
      # Note. VideoMAE (Facebook) adds positional encodinggs to the CLS token at
      # the encoder as well. VideoMAE (Tong et al) don't use a CLS token.
      # This implementation does not add positional embeddings to the CLS token,
      # as in the original image MAE of He et al.
      if input_shape is not None:
          input_shape = input_shape[:-1] + [
              self.decoder_config.hidden_size
          ]
      x_all = model_utils.add_positional_embeddings(
          x_all,
          self.positional_embedding_decoder,
          input_shape=input_shape,
          layer_name="posembed_decoder",
      )
  
      if self.classifier == "token":
          x_all = jnp.concatenate([cls_encoded, x_all], axis=1)
  
      x_decoded = vivit_model.Encoder(
          temporal_dims=temporal_dims,
          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, train=train)
  
      if self.classifier == "token":
          # Remove the CLS token for predicting reconstructions.
          x_decoded = x_decoded[:, 1:, :]
  
      x_prelogits = nn_layers.IdentityLayer(name="pre_logits")(
          x_decoded
      )
  
      x_logits = nn.Dense(
          self.num_classes,
          kernel_init=nn.initializers.zeros,
          name="output_projection",
      )(x_prelogits)
  
      return x_logits, token_mask
  
imports: from typing import Any, Optional
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import classification_model
  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.vivit import model
  from scenic.projects.vivit import model_utils
label_desc: Use image node_id 1 for calling the ViViTMaskedAutoencoder obj's underlying member method code with example usage: obj.__call__(inputs) and returns: x_logits,  token_mask, x_representation,  token_mask