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","."]
code: def __call__(
      self, inputs: jnp.ndarray, *, train: bool, debug: bool = False
  ):
      del debug
  
      # Shape is [batch, num_tokens, hidden_size]
      if self.modality == base_model.FeatureTargets.RGB:
          x_tokens, temporal_dims = vivit_model.temporal_encode(
              inputs,
              self.temporal_encoding_config,
              self.patches,
              self.hidden_size,
          )
      elif self.modality == base_model.FeatureTargets.SPECTROGRAM:
          x_tokens = model_utils.embed_2d_patch(
              inputs, self.patches, self.hidden_size
          )
          temporal_dims = None
      else:
          raise ValueError(f"Unknown modality {self.modality}")
  
      n_batch, n_tokens, hidden_dim = x_tokens.shape
      if self.modality == base_model.FeatureTargets.RGB:
          height = width = int(np.sqrt(n_tokens // temporal_dims))
          if height * width * temporal_dims != n_tokens:
              raise ValueError("Input is assumed to be square.")
  
      if (
          self.modality == base_model.FeatureTargets.SPECTROGRAM
          and self.positional_embedding
          not in ["learned_1d", "sinusoidal_1d"]
      ):
          raise ValueError(
              "Only 1d positional embdeddings are supported for spectograms."
          )
  
      # 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
      )
  
      x_tokens = self.add_modality_token(x_tokens)
  
      # If we want to add a class token, add it here.
      if self.classifier == "token":
          cls = self.param(
              "cls",
              nn.initializers.zeros,
              (1, 1, hidden_dim),
              x_tokens.dtype,
          )
          cls = jnp.tile(cls, [n_batch, 1, 1])
          x_tokens = jnp.concatenate([cls, x_tokens], axis=1)
  
      x_tokens_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_tokens, train=train)
  
      if self.classifier in ("token", "0"):
          x_tokens_encoded = x_tokens_encoded[:, 0]
      elif self.classifier in ("gap", "gmp", "gsp"):
          fn = {"gap": jnp.mean, "gmp": jnp.max, "gsp": jnp.sum}[
              self.classifier
          ]
          x_tokens_encoded = fn(x_tokens_encoded, axis=1)
          x_tokens_encoded = nn.LayerNorm(name="encoder_norm")(
              x_tokens_encoded
          )
      elif self.classifier in ("none"):
          return x_tokens_encoded
      else:
          raise ValueError(f"Unknown classifier {self.classifier}")
  
      if self.representation_size is not None:
          x_tokens_encoded = nn.Dense(
              self.representation_size, name="pre_logits"
          )(x_tokens_encoded)
          x_tokens_encoded = nn.tanh(x_tokens_encoded)
      else:
          x_tokens_encoded = nn_layers.IdentityLayer(name="pre_logits")(
              x_tokens_encoded
          )
  
      if self.freeze_backbone:
          x_tokens_encoded = jax.lax.stop_gradient(x_tokens_encoded)
  
          if self.use_batch_norm_after_encoder:
              x_tokens_encoded = nn.BatchNorm(
                  # Match PyTorch default and MAE PyTorch code
                  # https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/main_linprobe.py#L222
                  momentum=0.9,
                  epsilon=1e-6,
                  use_bias=False,
                  use_scale=False,
              )(x_tokens_encoded, use_running_average=not train)
  
      x_tokens_encoded = nn.Dense(
          self.num_classes,
          kernel_init=nn.initializers.zeros,
          name="output_projection",
      )(x_tokens_encoded)
      return x_tokens_encoded
  
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 ViViT obj's underlying member method code with example usage: obj.__call__(inputs) and returns: x_tokens_encoded, x_tokens_encoded