class_docstr: AdaTape Transformer model for Parity task.
  
  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.
  ac_config: Configuration of the adaptive computation.
  hidden_size: Size of the hidden state of the output of model's stem. 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: ["AdaTape","Transformer","model","for","Parity","task",".","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",".","ac_config",":","Configuration","of","the","adaptive","computation",".","hidden_size",":","Size","of","the","hidden","state","of","the","output","of","model","'s","stem",".","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.ac_config.bank_type == "input":
          x_input = x
          n, s, c = x.shape
          x = jnp.reshape(x, [n, 1, s * c])
          x = nn.Dense(self.hidden_size, name="embedding")(x)
          bank = nn.Dense(self.hidden_size, name="pre_bank")(x_input)
      else:
          raise NotImplementedError
  
      # If we want to add a class token, add it here.
      if self.classifier == "token":
          cls = self.param(
              "cls",
              nn.initializers.zeros,
              (1, 1, self.hidden_size),
              x.dtype,
          )
          cls = jnp.tile(cls, [n, 1, 1])
          x = jnp.concatenate([cls, x], axis=1)
  
      x, aux_output = adatape_layers.AdaTapeEncoder(
          mlp_dim=self.mlp_dim,
          num_layers=self.num_layers,
          num_heads=self.num_heads,
          ac_config=self.ac_config,
          dropout_rate=self.dropout_rate,
          attention_dropout_rate=self.attention_dropout_rate,
          stochastic_depth=self.stochastic_depth,
          dtype=self.dtype,
          name="Transformer",
      )(x, bank, 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)
  
      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, aux_output
  
imports: 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.common_lib import debug_utils
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.adatape import layers
  from scenic.projects.baselines import vit
label_desc: Use image node_id 1 for calling the AdaTapeParity obj's underlying member method code with example usage: obj.__call__(x) and returns: x,  aux_output