def __call__(self, x: Array, deterministic: bool) -> Array:
    """Applies MLP block update.

    Args:
      x: [..., input_dim].
      deterministic: don't apply dropout if true.

    Returns:
      Updated array x of same shape.
    """
    update = nn.gelu(self.mlp(x))
    update = self.dense(update)
    update = self.dropout(update, deterministic=deterministic)
    x = self.layer_norm(x + update)

    return x
Applies MLP block update.

    Args:
      x: [..., input_dim].
      deterministic: don't apply dropout if true.

    Returns:
      Updated array x of same shape.
    
language/mentionmemory/modules/mlp.py