class_docstr: Vision Transformer model.
  
  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.
  positional_embedding: The type of positional embeddings to add to the tokens
    at the beginning of the transformer encoder. Options are {learned_1d,
    sinusoidal_2d, none}.
  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.
  classifier: type of the classifier layer. Options are 'gap', 'gmp', 'gsp',
    'token'.
  dtype: JAX data type for activations.
class_docstr_tok: ["Vision","Transformer","model",".","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",".","positional_embedding",":","The","type","of","positional","embeddings","to","add","to","the","tokens","at","the","beginning","of","the","transformer","encoder",".","Options","are","{","learned_1d",",","sinusoidal_2d",",","none","}",".","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",".","classifier",":","type","of","the","classifier","layer",".","Options","are","'gap","'",",","'gmp","'",",","'gsp","'",",","'token","'",".","dtype",":","JAX","data","type","for","activations","."]
code: def __call__(
      self, x: jnp.ndarray, *, train: bool, debug: bool = False
  ):
      fh, fw = self.patches.size
      # Extracting patches and then embedding is in fact a single convolution.
      x = nn.Conv(
          self.hidden_size,
          (fh, fw),
          strides=(fh, fw),
          padding="VALID",
          name="embedding",
      )(x)
      n, h, w, c = x.shape
      x = jnp.reshape(x, [n, h * w, c])
  
      # Add posemb before adding extra token.
      x = x + get_posemb(
          self,
          self.positional_embedding,
          (h, w),
          c,
          "pos_embedding",
          x.dtype,
      )
  
      if self.classifier == "token":
          cls = self.param(
              "cls", nn.initializers.zeros, (1, 1, c), x.dtype
          )
          x = jnp.concatenate([jnp.tile(cls, [n, 1, 1]), x], axis=1)
  
      x = nn.Dropout(rate=self.dropout_rate)(x, not train)
  
      x = Encoder(
          depth=self.num_layers,
          mlp_dim=self.mlp_dim,
          num_heads=self.num_heads,
          dropout=self.dropout_rate,
          name="Transformer",
      )(x, deterministic=not train)
  
      if self.classifier == "map":
          x = MAPHead(num_heads=self.num_heads, mlp_dim=self.mlp_dim)(x)
      elif self.classifier == "gap":
          x = jnp.mean(x, axis=1)
      elif self.classifier in ("token", "0"):
          x = x[:, 0]
      else:
          raise ValueError(f"Unknown classifier {self.classifier}")
  
      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)
  
      return nn.Dense(
          self.num_classes,
          kernel_init=nn.initializers.zeros,
          name="output_projection",
      )(x)
  
imports: import functools
  from typing import Any, Optional
  from absl import logging
  import flax
  import flax.linen
  from flax.training import common_utils
  import jax.numpy
  import ml_collections
  import numpy
  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 import multilabel_classification_model
  from scenic.model_lib.layers import nn_layers
  import scipy
label_desc: Use image node_id 1 for calling the ViT obj's underlying member method code with example usage: obj.__call__(x) and returns: nn