class_docstr: Vision Transformer model for masked-autoencoding.
  
  The differences to `scenic.baselines.vit` are that:
  -- Remove masked tokens from the encoder.
  -- Process all tokens with the decoder.
  
  Attributes:
    num_classes: Number of output classes.
    mlp_dim: Dimension of the mlp on top of attention block.
    num_layers: Number of layers.
    num_heads: Number of self-attention heads.
    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 masking out the input tokens (with
      a learned mask token) during training.
    representation_size: Size of the representation layer in the model's head.
      if None, we skip the extra projection + tanh activation at the end.
    dropout_rate: Dropout rate.
    attention_dropout_rate: Dropout for attention heads.
    stochastic_depth: Probability of dropping out a layer during training.
    classifier: type of the classifier layer. Options are 'gap', 'gmp', 'gsp',
      'token'.
    dtype: JAX data type for activations.
class_docstr_tok: ["Vision","Transformer","model","for","masked-autoencoding",".","The","differences","to","`","scenic.baselines.vit","`","are","that",":","--","Remove","masked","tokens","from","the","encoder",".","--","Process","all","tokens","with","the","decoder",".","Attributes",":","num_classes",":","Number","of","output","classes",".","mlp_dim",":","Dimension","of","the","mlp","on","top","of","attention","block",".","num_layers",":","Number","of","layers",".","num_heads",":","Number","of","self-attention","heads",".","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","masking","out","the","input","tokens","(","with","a","learned","mask","token",")","during","training",".","representation_size",":","Size","of","the","representation","layer","in","the","model","'s","head",".","if","None",",","we","skip","the","extra","projection","+","tanh","activation","at","the","end",".","dropout_rate",":","Dropout","rate",".","attention_dropout_rate",":","Dropout","for","attention","heads",".","stochastic_depth",":","Probability","of","dropping","out","a","layer","during","training",".","classifier",":","type","of","the","classifier","layer",".","Options","are","'gap","'",",","'gmp","'",",","'gsp","'",",","'token","'",".","dtype",":","JAX","data","type","for","activations","."]
docstr: Forward pass of Vision Transformer.
docstr_tok: ["Forward","pass","of","Vision","Transformer","."]
code: def __call__(
      self, inputs: jnp.ndarray, *, train: bool, debug: bool = False
  ):
      """Forward pass of Vision Transformer."""
  
      # Extracting patches and embed via a convolution.
      fh, fw = self.patches.size
      x_tokens = nn.Conv(
          self.hidden_size,
          (fh, fw),
          strides=(fh, fw),
          padding="VALID",
          name="embedding",
      )(inputs)
      n_batch, height, width, channels = x_tokens.shape
      n_tokens = height * width
      x_tokens = jnp.reshape(x_tokens, [n_batch, n_tokens, channels])
  
      # Add positional encodings.
      x_tokens = model_utils.add_positional_embeddings(
          x_tokens,
          self.positional_embedding,
          [n_batch, height, width, channels],
      )
  
      if train:
          # Generate mask indices.
          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")
          )
  
          # 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, channels),
              x_unmasked.dtype,
          )
          cls = jnp.tile(cls, [n_batch, 1, 1])
          x_unmasked = jnp.concatenate([cls, x_unmasked], axis=1)
  
      x_unmasked_encoded = vit.Encoder(
          mlp_dim=self.mlp_dim,
          num_layers=self.num_layers,
          num_heads=self.num_heads,
          dropout_rate=self.dropout_rate,
          attention_dropout_rate=self.attention_dropout_rate,
          stochastic_depth=self.stochastic_depth,
          dtype=self.dtype,
          positional_embedding="none",  # Has already been added.
          name="Transformer",
      )(x_unmasked, train=train)
  
      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,
          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)
  
      # Add positional encodings to the decoder.
      x_all = model_utils.add_positional_embeddings(
          x_all,
          self.positional_embedding_decoder,
          [n_batch, height, width, self.decoder_config.hidden_size],
      )
  
      if self.classifier == "token":
          x_all = jnp.concatenate([cls_encoded, x_all], axis=1)
  
      x_decoded = vit.Encoder(
          mlp_dim=self.decoder_config.mlp_dim,
          num_layers=self.decoder_config.num_layers,
          num_heads=self.decoder_config.num_heads,
          dropout_rate=self.decoder_config.dropout_rate,
          attention_dropout_rate=self.decoder_config.attention_dropout_rate,
          stochastic_depth=self.decoder_config.stochastic_depth,
          dtype=self.dtype,
          positional_embedding="none",  # Has already been added.
          name="Decoder",
      )(x_all, train=train)
  
      if self.classifier == "token":
          # Remove the CLS token for predicting reconstructions.
          x_decoded = x_decoded[:, 1:, :]
  
      # Predict pixel reconstructions.
      if self.representation_size is not None:
          x_prelogits = nn.Dense(
              self.representation_size, name="pre_logits"
          )(x_decoded)
          x_prelogits = nn.tanh(x_prelogits)
      else:
          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, Callable, Optional, Sequence
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import multilabel_classification_model
  from scenic.model_lib.layers import attention_layers
  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.baselines import vit
label_desc: Use image node_id 1 for calling the ViTMaskedAutoencoder obj's underlying member method code with example usage: obj.__call__(inputs) and returns: x_logits,  token_mask, x_representation,  token_mask