class_docstr: Stem for BERT.
  
  Attributes:
    vocab_size: Size of words/tokens vocabulary.
    type_vocab_size: Size of type vocabulary.
    hidden_size: Size of the hidden state of the output of model's stem.
    max_position_embeddings: The maximum sequence length that this model might
      ever be used with.
    embedding_width: Size of embedding
    dropout_rate: Dropout rate.
    dtype: JAX data type for activations.
class_docstr_tok: ["Stem","for","BERT",".","Attributes",":","vocab_size",":","Size","of","words\/tokens","vocabulary",".","type_vocab_size",":","Size","of","type","vocabulary",".","hidden_size",":","Size","of","the","hidden","state","of","the","output","of","model","'s","stem",".","max_position_embeddings",":","The","maximum","sequence","length","that","this","model","might","ever","be","used","with",".","embedding_width",":","Size","of","embedding","dropout_rate",":","Dropout","rate",".","dtype",":","JAX","data","type","for","activations","."]
code: def __call__(
      self,
      input_word_ids: jnp.ndarray,
      input_type_ids: jnp.ndarray,
      input_mask: jnp.ndarray,
      *,
      train: bool,
  ) -> Union[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray]]:
      if input_word_ids.ndim != 2:
          raise ValueError(
              "Input_word_ids should be of shape `[bs, l]` but it is "
              f"{input_word_ids.shape}."
          )
      if input_type_ids.ndim != 2:
          raise ValueError(
              "Input_type_ids should be of shape `[bs, l]` but it is "
              f"{input_type_ids.shape}."
          )
      if input_mask.ndim != 2:
          raise ValueError(
              "Input_mask should be of shape `[bs, l]` but it is "
              f"{input_mask.shape}."
          )
  
      embedding_width = (
          self.embedding_width
          if self.embedding_width
          else self.hidden_size
      )
  
      word_embedding_layer = nn.Embed(
          num_embeddings=self.vocab_size,
          features=embedding_width,
          embedding_init=bert_truncated_normal_initializer(),
          name="word_embedding",
      )
      x = word_embedding_layer(input_word_ids)
      x = x + nn.Embed(
          num_embeddings=self.type_vocab_size,
          features=embedding_width,
          embedding_init=bert_truncated_normal_initializer(),
          name="type_embedding",
      )(input_type_ids)
      # NOTE: CLS token is added during pre-processing in the tokenizer.
      x = AddPositionEmbs(
          max_len=self.max_position_embeddings,
          posemb_init=bert_truncated_normal_initializer(),
          name="posembed_input",
      )(x)
      x = nn.LayerNorm(dtype=self.dtype, name="embedding_norm")(x)
      x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=not train)
      if embedding_width != self.hidden_size:
          x = nn.Dense(
              self.hidden_size,
              kernel_init=bert_truncated_normal_initializer(),
              name="embedding_projection",
          )(x)
      return x, word_embedding_layer.embedding
  
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 Stem obj's underlying member method code with example usage: obj.__call__(input_word_ids, input_type_ids, input_mask) and returns: x,  word_embedding_layer