class_docstr: Tokenizer for 2D inputs (e.g., images).
  
  Attributes:
    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.
    mlp_dim: Dimension of the mlp on top of attention block.
    num_layers: Number of layers.
    num_heads: Number of self-attention heads.
    dropout_rate: Dropout rate.
    attention_dropout_rate: Dropout for attention heads.
    add_cls_token: Whether to add CLS token or not.
    dtype: JAX data type for activations.
class_docstr_tok: ["Tokenizer","for","2D","inputs","(","e.g.",",","images",")",".","Attributes",":","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",".","mlp_dim",":","Dimension","of","the","mlp","on","top","of","attention","block",".","num_layers",":","Number","of","layers",".","num_heads",":","Number","of","self-attention","heads",".","dropout_rate",":","Dropout","rate",".","attention_dropout_rate",":","Dropout","for","attention","heads",".","add_cls_token",":","Whether","to","add","CLS","token","or","not",".","dtype",":","JAX","data","type","for","activations","."]
code: def __call__(
      self,
      x: jnp.ndarray,
      *,
      train: bool,
      dataset: str,
      stochastic_droplayer_rate: Optional[float],
  ) -> jnp.ndarray:
      if x.ndim != 4:
          raise ValueError(
              f"Input shape should be `[bs, h, w, c]` but it is {x.shape}."
          )
      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 self.add_cls_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.AddPositionEmbs(
          posemb_init=nn.initializers.normal(stddev=0.02),  # from BERT.
          name="posembed_input",
      )(x)
  
      x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=not train)
  
      for lyr in range(self.num_layers):
          droplayer_p = get_droplayer_p(
              lyr, self.num_layers, stochastic_droplayer_rate
          )
          x = Encoder1DBlock(
              mlp_dim=self.mlp_dim,
              num_heads=self.num_heads,
              dropout_rate=self.dropout_rate,
              attention_dropout_rate=self.attention_dropout_rate,
              name=f"encoderblock_{lyr}",
              dtype=jax.dtypes.canonicalize_dtype(self.dtype),
          )(
              x,
              deterministic=not train,
              droplayer_p=droplayer_p,
              dataset=dataset,
          )
  
      return x
  
imports: import functools
  from typing import Any, Callable, Optional, Sequence, Tuple, Union
  import flax.linen
  import jax
  from jax.nn import initializers
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
  from scenic.projects.polyvit import polyvit_base_model
  from scenic.projects.vivit import model
label_desc: Use image node_id 1 for calling the Tokenizer2D obj's underlying member method code with example usage: obj.__call__(x) and returns: x