class_docstr: Vision Transformer model.
  
  This differs from scenic.projects.baselines.vit in that
  -- Positional embeddings are added before the transformer block. This makes
     it easier to load MAE-pretrained checkpoints.
  -- The CLS token is added after the positional embeddings. This follows MAE.
  -- Add support for linear evaluation by adding a stop-gradient.
  
  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'.
  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).
  dtype: JAX data type for activations.
class_docstr_tok: ["Vision","Transformer","model",".","This","differs","from","scenic.projects.baselines.vit","in","that","--","Positional","embeddings","are","added","before","the","transformer","block",".","This","makes","it","easier","to","load","MAE-pretrained","checkpoints",".","--","The","CLS","token","is","added","after","the","positional","embeddings",".","This","follows","MAE",".","--","Add","support","for","linear","evaluation","by","adding","a","stop-gradient",".","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","'",".","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",")",".","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 positional embeddings to tokens.
      if self.positional_embedding == "learned_1d":
          x = vit.AddPositionEmbs(
              posemb_init=nn.initializers.normal(
                  stddev=0.02
              ),  # from BERT.
              name="posembed_input",
          )(x)
      elif self.positional_embedding == "sinusoidal_1d":
          x = attention_layers.Add1DPositionEmbedding(posemb_init=None)(
              x
          )
      elif self.positional_embedding == "sinusoidal_2d":
          x_reshape = x.reshape([n, h, w, c])
          x = attention_layers.AddFixedSinCosPositionEmbedding()(
              x_reshape
          )
          x = jnp.reshape(x, [n, h * w, c])
      elif self.positional_embedding == "none":
          pass
      else:
          raise ValueError(
              "Unknown positional embedding: "
              f"{self.positional_embedding}"
          )
  
      # 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 = vit.Encoder(
          mlp_dim=self.mlp_dim,
          num_layers=self.num_layers,
          num_heads=self.num_heads,
          positional_embedding="none",
          dropout_rate=self.dropout_rate,
          attention_dropout_rate=self.attention_dropout_rate,
          stochastic_depth=self.stochastic_depth,
          dtype=self.dtype,
          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 == "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.freeze_backbone:
          x = jax.lax.stop_gradient(x)
  
          if self.use_batch_norm_after_encoder:
              x = 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, use_running_average=not train)
  
      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
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import multilabel_classification_model
  from scenic.model_lib.layers import attention_layers
  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.baselines import vit
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