class_docstr: Vision Video Transformer model baseline for transfer learning.
  
    The differences to the scenic.project.vivit.model.ViViT are that:
    -- Posibility of freezing the backbone.
    -- The positional embedding is added before the encoder.
    -- 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.
    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.
    representation_size: Size of the representation layer in the model's head.
      if None, we skip the extra projection + tanh activation at the end.
    classifier: type of the classifier layer. Options are 'gap', 'gmp', 'gsp',
      'token'.
    freeze_backbone: If True, add a stop-gradient before the final
      classifier to only evaluate linear evaluation performance.
    use_batch_norm_after_encoder: Only applies when the backbone is frozen.
      In this case, an additional batch normalisation layer is applied before
      the linear classifier. This was done in MAE
      (https://arxiv.org/abs/2111.06377).
    positional_embedding: The type of positional embeddings to add to the
      tokens at the beginning of the transformer encoder. Options are
      {learned_1d, sinusoidal_3d, none}.
    normalise_encoder_output: If true, layer normalisation is applied to the
      output of the transformer encoder. This is typically not done when
      not using the token classifier.
    dtype: JAX data type for activations.
class_docstr_tok: ["Vision","Video","Transformer","model","baseline","for","transfer","learning",".","The","differences","to","the","scenic.project.vivit.model.ViViT","are","that",":","--","Posibility","of","freezing","the","backbone",".","--","The","positional","embedding","is","added","before","the","encoder",".","--","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",".","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",".","representation_size",":","Size","of","the","representation","layer","in","the","model","'s","head",".","if","None",",","we","skip","the","extra","projection","+","tanh","activation","at","the","end",".","classifier",":","type","of","the","classifier","layer",".","Options","are","'gap","'",",","'gmp","'",",","'gsp","'",",","'token","'",".","freeze_backbone",":","If","True",",","add","a","stop-gradient","before","the","final","classifier","to","only","evaluate","linear","evaluation","performance",".","use_batch_norm_after_encoder",":","Only","applies","when","the","backbone","is","frozen",".","In","this","case",",","an","additional","batch","normalisation","layer","is","applied","before","the","linear","classifier",".","This","was","done","in","MAE","(","https",":","\/\/arxiv.org\/abs\/2111.06377",")",".","positional_embedding",":","The","type","of","positional","embeddings","to","add","to","the","tokens","at","the","beginning","of","the","transformer","encoder",".","Options","are","{","learned_1d",",","sinusoidal_3d",",","none","}",".","normalise_encoder_output",":","If","true",",","layer","normalisation","is","applied","to","the","output","of","the","transformer","encoder",".","This","is","typically","not","done","when","not","using","the","token","classifier",".","dtype",":","JAX","data","type","for","activations","."]
docstr: Add modality learned tokens.
docstr_tok: ["Add","modality","learned","tokens","."]
code: def add_modality_token(
      self, x_tokens: jnp.ndarray, name: str = "Encoder"
  ) -> jnp.ndarray:
      """Add modality learned tokens."""
      if not self.use_modality_tokens:
          return x_tokens
  
      modality_token = self.param(
          f"{name}_modality_token_{self.modality}",
          nn.initializers.zeros,
          (1, 1, x_tokens.shape[-1]),
      )
      x_tokens = x_tokens + modality_token
      return x_tokens
  
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 2 for calling the ViViT obj's underlying member method code with example usage: obj.add_modality_token(x_tokens, name) and returns: x_tokens, x_tokens