def __call__(
      self,
      encoding: Array,
      attention_mask: Array,
      deterministic: bool,
  ) -> Array:
    """Transformer layer forward.

    Args:
      encoding: [bsz, seq_len, model_dim] model state.
      attention_mask: [bsz, seq_len].
      deterministic: if true, do not apply dropout.

    Returns:
      updated encoding
    """

    encoding = self.attention_block(
        encoding=encoding,
        attention_mask=attention_mask,
        deterministic=deterministic)

    encoding = self.mlp_block(x=encoding, deterministic=deterministic)

    return encoding
Transformer layer forward.

    Args:
      encoding: [bsz, seq_len, model_dim] model state.
      attention_mask: [bsz, seq_len].
      deterministic: if true, do not apply dropout.

    Returns:
      updated encoding
    
language/mentionmemory/modules/transformer.py