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.
  hidden_size: Size of the hidden state of the output of model's stem.
  resnet: Configuration of the ResNet block in the stem of the model.
  patches: Configuration of the patches extracted in the stem of the model.
  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",".","hidden_size",":","Size","of","the","hidden","state","of","the","output","of","model","'s","stem",".","resnet",":","Configuration","of","the","ResNet","block","in","the","stem","of","the","model",".","patches",":","Configuration","of","the","patches","extracted","in","the","stem","of","the","model",".","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
  ):
      if self.resnet is not None:
          width = int(64 * self.resnet.width_factor)
          # Root block.
          x = bit_resnet.StdConv(
              width,
              (7, 7),
              strides=(2, 2),
              use_bias=False,
              name="conv_root",
          )(x)
          x = nn.GroupNorm(name="gn_root")(x)
          x = nn.relu(x)
          x = nn.max_pool(x, (3, 3), strides=(2, 2), padding="SAME")
          # ResNet stages.
          if self.resnet.num_layers:
              x = bit_resnet.ResNetStage(
                  block_size=self.resnet.num_layers[0],
                  nout=width,
                  first_stride=(1, 1),
                  name="block1",
              )(x)
              for i, block_size in enumerate(
                  self.resnet.num_layers[1:], 1
              ):
                  x = bit_resnet.ResNetStage(
                      block_size=block_size,
                      nout=width * 2**i,
                      first_stride=(2, 2),
                      name=f"block{i + 1}",
                  )(x)
      fh, fw = self.patches.size
      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 = vit.Encoder(
          mlp_dim=self.mlp_dim,
          num_layers=self.num_layers,
          num_heads=self.num_heads,
          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)
  
      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)
      x = nn.Dense(
          self.num_classes,
          kernel_init=nn.initializers.zeros,
          name="output_projection",
      )(x)
      return x
  
file_path: scenic/scenic/projects/baselines/hybrid_vit.py
imports: from typing import Any, Optional
  import flax.linen
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import bit_resnet
  from scenic.projects.baselines import vit
label_desc: Use image node_id 1 for calling the HybridViT obj's underlying member method code with example usage: obj.__call__(x) and returns: x