class_docstr: BERT encoder.
  
  Attributes:
  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.
  pre_norm: Whether to use PreLN in encoder layers, otherwise PostLN.
  dtype: JAX data type for activations.
class_docstr_tok: ["BERT","encoder",".","Attributes",":","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",".","pre_norm",":","Whether","to","use","PreLN","in","encoder","layers",",","otherwise","PostLN",".","dtype",":","JAX","data","type","for","activations","."]
code: def __call__(
      self, x: jnp.ndarray, input_mask: jnp.ndarray, *, train: bool
  ):
      for lyr in range(self.num_layers):
          x = Encoder1DBlock(
              mlp_dim=self.mlp_dim,
              num_heads=self.num_heads,
              dropout_rate=self.dropout_rate,
              attention_dropout_rate=self.attention_dropout_rate,
              pre_norm=self.pre_norm,
              name=f"encoderblock_{lyr}",
              dtype=jax.dtypes.canonicalize_dtype(self.dtype),
          )(x, input_mask=input_mask, deterministic=not train)
  
      if self.pre_norm:
          x = nn.LayerNorm(name="encoder_norm")(x)
      return x
  
imports: from typing import Any, Callable, Optional, Sequence, Tuple, Union
  import flax.linen
  import jax
  from jax.nn import initializers
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 1 for calling the BERTEncoder obj's underlying member method code with example usage: obj.__call__(x, input_mask) and returns: x