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', 'none'.
  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])
  
      # If we want to add a class token, add it here.
      if self.classifier == "token":
          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(
          mlp_dim=self.mlp_dim,
          num_layers=self.num_layers,
          num_heads=self.num_heads,
          positional_embedding=self.positional_embedding,
          dropout_rate=self.dropout_rate,
          attention_dropout_rate=self.attention_dropout_rate,
          stochastic_depth=self.stochastic_depth,
          dtype=self.dtype,
          has_cls_token=self.classifier == "token",
          name="Transformer",
      )(x, train=train)
  
      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=1)
      elif self.classifier == "map":
          x = MAPHead(
              num_heads=self.num_heads,
              mlp_dim=self.mlp_dim,
              dtype=self.dtype,
          )(x)
      elif self.classifier == "none":
          pass
      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)
  
      if self.num_classes > 0:
          # If self.num_classes <= 0, we just return the backbone features.
          x = nn.Dense(
              self.num_classes,
              kernel_init=nn.initializers.zeros,
              name="output_projection",
          )(x)
      return x
  
imports: from typing import Any, Callable, Optional, Sequence
  from absl import logging
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  import scipy
  from tensorflow.io import gfile
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: x