class_docstr: XViT model.
  
  Attributes:
    num_outputs: number of classes.
    mlp_dim: Dimension of the MLP on top of attention block.
    num_layers: Number of layers.
    attention_configs: Configurations passed to the self-attention.
    attention_fn: Self-attention function used in the model.
    patches: Configuration of the patches extracted in the stem of the model.
    hidden_size: Size of the hidden dimension on the stem of the model.
    fast_vit: Configurations of the fast_vit (omnidirectional attention).
    representation_size: Size of the final representation.
    dropout_rate: Dropout rate.
    attention_dropout_rate: Dropout rate for attention heads.
    classifier:  Type of the classifier.
class_docstr_tok: ["XViT","model",".","Attributes",":","num_outputs",":","number","of","classes",".","mlp_dim",":","Dimension","of","the","MLP","on","top","of","attention","block",".","num_layers",":","Number","of","layers",".","attention_configs",":","Configurations","passed","to","the","self-attention",".","attention_fn",":","Self-attention","function","used","in","the","model",".","patches",":","Configuration","of","the","patches","extracted","in","the","stem","of","the","model",".","hidden_size",":","Size","of","the","hidden","dimension","on","the","stem","of","the","model",".","fast_vit",":","Configurations","of","the","fast_vit","(","omnidirectional","attention",")",".","representation_size",":","Size","of","the","final","representation",".","dropout_rate",":","Dropout","rate",".","attention_dropout_rate",":","Dropout","rate","for","attention","heads",".","classifier",":","Type","of","the","classifier","."]
docstr: X-ViT model.
  
  Args:
    inputs: Input data.
    train: If it is training.
    debug: If we are running the model in the debug mode.
  
  Returns:
    Output of the model.
docstr_tok: ["X-ViT","model",".","Args",":","inputs",":","Input","data",".","train",":","If","it","is","training",".","debug",":","If","we","are","running","the","model","in","the","debug","mode",".","Returns",":","Output","of","the","model","."]
code: def __call__(
      self,
      inputs: jnp.ndarray,
      *,
      train: bool,
      debug: Optional[bool] = False,
  ) -> jnp.ndarray:
      """X-ViT model.
  
      Args:
        inputs: Input data.
        train: If it is training.
        debug: If we are running the model in the debug mode.
  
      Returns:
        Output of the model.
      """
      _, height, width, _ = inputs.shape
      patch_height, patch_width = self.patches.size
      grid_height, grid_width = (
          height // patch_height,
          width // patch_width,
      )
      patch_stride_height, patch_stride_width = self.patches.get(
          "strides", self.patches.size
      )
      x = nn.Conv(
          self.hidden_size,
          (patch_height, patch_width),
          strides=(patch_stride_height, patch_stride_width),
          padding="VALID",
          name="embedding",
      )(inputs)
  
      transformer_encoder_type = self.transformer_encoder_configs.type
      if transformer_encoder_type in ["global", "pyramid"]:
          bs, h, w, c = x.shape
          x = jnp.reshape(x, [bs, h * w, c])
          if self.classifier == "token":
              # Only when we do flattening here, we can add the extra CLS token, in
              #  other cases, we use the first token as CLS token.
              cls = self.param(
                  "cls", nn.initializers.zeros, (1, 1, c), x.dtype
              )
              cls = jnp.tile(cls, [bs, 1, 1])
              x = jnp.concatenate([cls, x], axis=1)
  
      if transformer_encoder_type == "global":
          x = Encoder1D(
              mlp_dim=self.mlp_dim,
              num_layers=self.num_layers,
              attention_configs=self.attention_configs,
              attention_fn=self.attention_fn,
              dropout_rate=self.dropout_rate,
              attention_dropout_rate=self.attention_dropout_rate,
              name="Transformer",
          )(x, train=train)
          cls_token = x[:, 0]
  
      elif transformer_encoder_type == "pyramid":
          if self.transformer_encoder_configs["pyramid_mode"] in [
              "2d_pool",
              "cnn",
          ]:
              assert grid_height == grid_width, (
                  "For now, only square grid is supported for pyramid with 2d "
                  "pooling or cnn pooling."
              )
          x = Encoder1DPyramid(
              mlp_dim=self.mlp_dim,
              num_layers=self.fnum_layers,
              attention_configs=self.attention_configs,
              attention_fn=self.attention_fn,
              transformer_encoder_configs=self.transformer_encoder_configs,
              dropout_rate=self.dropout_rate,
              attention_dropout_rate=self.attention_dropout_rate,
              classifier=self.classifier,
              name="Transformer",
          )(x, train=train)
          cls_token = x[:, 0]
  
      elif transformer_encoder_type in [
          "axial",
          "axial_attention",
          "grid",
          "grid_attention",
      ]:
          if transformer_encoder_type in ["grid", "grid_attention"]:
              # First put patches of patches in rows. (we already projected pixels to
              # patches in the stem and at this level, the input tokens are patches).
              # TODO(dehghani): Check if nn_ops.extract_image_patches is faster.
              pp_size = (
                  self.transformer_encoder_configs.patches_of_patches_size
              )
              x = nn_ops.extract_patches(
                  lhs=x, rhs_shape=pp_size, strides=pp_size
              )
              # TODO(dehghani): Check if we can output a 4D tensor directly and get
              #  rid of this reshape.
              bs, ph, pw, h, w, c = x.shape
              x = x.reshape(bs, ph * pw, h * w, c)
              # Now, we simply need to run axial/axial_attention to get
              # inter-patch and intra-patch attention.
              if transformer_encoder_type == "grid":
                  transformer_encoder_type = "axial"
              elif transformer_encoder_type == "grid_attention":
                  transformer_encoder_type = "axial_attention"
  
          x = EncoderAxial(
              mlp_dim=self.mlp_dim,
              num_layers=self.num_layers,
              transformer_encoder_type=transformer_encoder_type,
              attention_configs=self.attention_configs,
              attention_fn=self.attention_fn,
              dropout_rate=self.dropout_rate,
              attention_dropout_rate=self.attention_dropout_rate,
              name="Transformer",
          )(x, train=train)
          cls_token = x[:, 0, 0]
  
      else:
          raise ValueError(
              f"Transformer encoder type {transformer_encoder_type} is not defined!"
          )
  
      if self.classifier in ("token", "0"):
          x = cls_token
  
      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 = jnp.mean(x, axis=list(range(1, x.ndim - 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_outputs,
          kernel_init=nn.initializers.zeros,
          name="output_projection",
      )(x)
      return x
  
imports: from typing import Any, Optional
  import flax.linen
  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
  from scenic.model_lib.layers import nn_ops
  from scenic.projects.baselines import vit
  from scenic.projects.fast_vit import model_utils
label_desc: Use image node_id 1 for calling the XViT obj's underlying member method code with example usage: obj.__call__(inputs) and returns: x