def __call__(self, encoding: Array, attention_mask: Array,
               deterministic: bool) -> Array:
    """Self attention 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.
    """

    attention_mask = nn.make_attention_mask(attention_mask, attention_mask)
    update = self.attention_layer(
        inputs_q=encoding, mask=attention_mask, deterministic=deterministic)
    update = self.dropout(update, deterministic=deterministic)
    encoding = self.layer_norm(encoding + update)

    return encoding
Self attention 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/attention.py