class_docstr: Vision Transformer model for Video.
  
  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.
    representation_size: Size of the representation layer in the model's head.
      if None, we skip the extra projection + tanh activation at the end.
    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'.
    return_prelogits: If true, return the final representation of the network
      before the classification head. Useful when using features for a
      downstream task.
    return_preclassifier: If true, return the representation after the
      transformer encoder. Useful if using this as the backbone stem as part
      of a bigger architecture.
    dtype: JAX data type for activations.
class_docstr_tok: ["Vision","Transformer","model","for","Video",".","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",".","representation_size",":","Size","of","the","representation","layer","in","the","model","'s","head",".","if","None",",","we","skip","the","extra","projection","+","tanh","activation","at","the","end",".","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","'",".","return_prelogits",":","If","true",",","return","the","final","representation","of","the","network","before","the","classification","head",".","Useful","when","using","features","for","a","downstream","task",".","return_preclassifier",":","If","true",",","return","the","representation","after","the","transformer","encoder",".","Useful","if","using","this","as","the","backbone","stem","as","part","of","a","bigger","architecture",".","dtype",":","JAX","data","type","for","activations","."]
code: def __call__(
      self, x: jnp.ndarray, *, train: bool, debug: bool = False
  ):
      assert self.classifier in ["token", "0", "gap", "gmp", "gsp"]
      attention_type = self.attention_config.get("type", "spacetime")
      if attention_type in [
          "factorized_transformer_block",
          "factorized_self_attention_block",
          "factorized_dot_product_attention",
      ]:
          assert self.classifier not in ["token", "0"], (
              "For factorized_transformer_block, factorized_self_attention_block"
              "and factorized_dot_product_attention, the token classifier is not"
              "implemented."
          )
  
      x, temporal_dims = temporal_encode(
          x,
          self.temporal_encoding_config,
          self.patches,
          self.hidden_size,
      )
  
      # If we want to add a class token, add it here.
      if self.classifier in ["token"]:
          n, _, c = x.shape
          cls = self.param(
              "cls", nn.initializers.zeros, (1, 1, c), x.dtype
          )
          cls = jnp.tile(cls, [n, 1, 1])
          x = jnp.concatenate([cls, x], axis=1)
  
      x = 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,
          name="Transformer",
      )(x, train=train)
  
      if self.return_preclassifier:
          return x
  
      if self.classifier in ["token", "0"]:
          x = x[:, 0]
      elif self.classifier in ("gap", "gmp", "gsp"):
          fn = {"gap": jnp.mean, "gmp": jnp.max, "gsp": jnp.sum}[
              self.classifier
          ]
          x = fn(x, axis=list(range(1, x.ndim - 1)))
  
      if self.representation_size is not None:
          x = nn.Dense(self.representation_size, name="pre_logits")(x)
          x = nn.tanh(x)
      else:
          x = nn_layers.IdentityLayer(name="pre_logits")(x)
  
      if self.return_prelogits:
          return x
      else:
          x = nn.Dense(
              self.num_classes,
              kernel_init=nn.initializers.zeros,
              name="output_projection",
          )(x)
          return x
  
imports: import functools
  from typing import Any, Optional, Callable, Sequence
  from absl import logging
  import flax.linen
  from flax.linen.linear import default_kernel_init
  from immutabledict import immutabledict
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.common_lib import video_utils
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
  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__(x) and returns: x, x, x