class_docstr: Vision Transformer.
  
  Attributes:
    patch_size: The size of the patches to embed.
    features: Number of features.
    num_layers: Number of transformer blocks (self-attn + MLP).
    num_heads: Number of attention heads.
    out_features: Number of output features. If None, return transformer output.
    stochastic_droplayer_rate: Stochastic depth rate.
class_docstr_tok: ["Vision","Transformer",".","Attributes",":","patch_size",":","The","size","of","the","patches","to","embed",".","features",":","Number","of","features",".","num_layers",":","Number","of","transformer","blocks","(","self-attn","+","MLP",")",".","num_heads",":","Number","of","attention","heads",".","out_features",":","Number","of","output","features",".","If","None",",","return","transformer","output",".","stochastic_droplayer_rate",":","Stochastic","depth","rate","."]
code: def __call__(
      self,
      x: jnp.ndarray,
      attn_mask: Optional[jnp.ndarray] = None,
      *,
      deterministic: bool = True,
  ) -> jnp.ndarray:
      x = nn.Conv(
          self.features,
          kernel_size=(self.patch_size, self.patch_size),
          strides=(self.patch_size, self.patch_size),
          use_bias=False,
          name="conv1",
      )(x)
      x = x.reshape(x.shape[0], -1, x.shape[-1])
      scale = 1.0 / jnp.sqrt(self.features)
      class_embedding = self.param(
          "class_embedding",
          jax.nn.initializers.normal(stddev=scale),
          (self.features,),
      )
      x = jnp.concatenate(
          (
              jnp.tile(
                  class_embedding[None, None, :], (x.shape[0], 1, 1)
              ),
              x,
          ),
          axis=1,
      )
      positional_embedding = self.param(
          "positional_embedding",
          jax.nn.initializers.normal(stddev=scale),
          (x.shape[1], self.features),
      )
      x = x + positional_embedding[None]
  
      x = LayerNorm(name="ln_pre")(x)
      x = feature_map = Transformer(
          features=self.features,
          num_layers=self.num_layers,
          num_heads=self.num_heads,
          stochastic_droplayer_rate=self.stochastic_droplayer_rate,
          name="transformer",
      )(x, deterministic=deterministic)
  
      if self.out_features is not None:
          x = LayerNorm(name="ln_post")(x[:, 0])
          x = nn.Dense(self.out_features, use_bias=False, name="proj")(
              x
          )
      else:
          x = LayerNorm(name="ln_post")(x)
  
      return x, feature_map
  
imports: import functools
  from typing import Sequence, Optional, Union, Tuple
  import flax.linen
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling the VisionTransformer obj's underlying member method code with example usage: obj.__call__(x, attn_mask) and returns: x,  feature_map