def __call__(self, input_ids, pos_ids=None, token_type_ids=None):
    """Construct the Transformer embeddings.

    Args:
      input_ids: <int32>[batch_size, len].
      pos_ids: <int32>[batch_size, len].
      token_type_ids: <int32>[batch_size, len].

    Returns:
      embs: <float32>[batch_size, len, hidden_size].
    """
    batch_size, seq_length = input_ids.shape
    if pos_ids is None:
      pos_ids = jnp.tile(
          jnp.arange(seq_length, dtype=jnp.int32)[None, :], (batch_size, 1))
      if self.decode:
        # cache position index for tracking decoding position.
        is_initialized = self.has_variable('cache', 'cache_index')
        cache_index = self.variable('cache', 'cache_index',
                                    lambda: jnp.array(0, dtype=jnp.uint32))
        if is_initialized:
          i = cache_index.value
          cache_index.value = i + 1
          pos_ids = jnp.full(input_ids.shape, i, dtype=jnp.int32)
    if token_type_ids is None:
      token_type_ids = jnp.zeros_like(input_ids)
    input_embs = nn.Embed(
        self.vocab_size,
        self.hidden_size,
        embedding_init=self.embedding_init,
        name='word_embeddings')(
            input_ids)
    position_embs = nn.Embed(
        self.max_position_embeddings,
        self.hidden_size,
        embedding_init=self.embedding_init,
        name='position_embeddings')(
            pos_ids)
    token_type_embs = nn.Embed(
        self.type_vocab_size,
        self.hidden_size,
        embedding_init=self.embedding_init,
        name='token_type_embeddings')(
            token_type_ids)
    embs = input_embs + position_embs + token_type_embs
    embs = nn.LayerNorm(epsilon=self.layer_norm_eps, name='LayerNorm')(embs)
    embs = nn.Dropout(rate=self.hidden_dropout_rate)(
        embs, deterministic=self.deterministic)
    return embs
Construct the Transformer embeddings.

    Args:
      input_ids: <int32>[batch_size, len].
      pos_ids: <int32>[batch_size, len].
      token_type_ids: <int32>[batch_size, len].

    Returns:
      embs: <float32>[batch_size, len, hidden_size].
    
language/gscan/xattn_model/model/layers.py