class_docstr: PolyViT encoder.
  
  Attributes:
  mlp_dim: Dimension of the mlp on top of attention block.
  num_tokenizer_layers: Number of tokenizer layers.
  num_layers: Number of layers.
  num_heads: Number of self-attention heads.
  dropout_rate: Dropout rate.
  attention_dropout_rate: Dropout for attention heads.
  dtype: JAX data type for activations.
class_docstr_tok: ["PolyViT","encoder",".","Attributes",":","mlp_dim",":","Dimension","of","the","mlp","on","top","of","attention","block",".","num_tokenizer_layers",":","Number","of","tokenizer","layers",".","num_layers",":","Number","of","layers",".","num_heads",":","Number","of","self-attention","heads",".","dropout_rate",":","Dropout","rate",".","attention_dropout_rate",":","Dropout","for","attention","heads",".","dtype",":","JAX","data","type","for","activations","."]
code: def __call__(
      self,
      x: jnp.ndarray,
      *,
      train: bool,
      dataset: str,
      stochastic_droplayer_rate: Optional[float],
  ):
      for lyr in range(self.num_tokenizer_layers, 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,
          )
  
      x = nn.LayerNorm(name="encoder_norm")(x)
  
      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 PolyViTEncoder obj's underlying member method code with example usage: obj.__call__(x) and returns: x