def __call__(self, hidden, x):
    """Apply the Transformer output layer.

    Args:
      hidden: <float32>[batch_size, len, dim].
      x: <float32>[batch_size, len, dim].

    Returns:
      output: <float32>[batch_size, len, hidden_size].
    """
    hidden = nn.Dense(
        self.hidden_size,
        kernel_init=self.kernel_init,
        bias_init=self.bias_init,
        name='dense')(
            hidden)
    hidden = nn.Dropout(rate=self.hidden_dropout_rate)(
        hidden, deterministic=self.deterministic)
    hidden = nn.LayerNorm(
        epsilon=self.layer_norm_eps, name='LayerNorm')(
            hidden + x)
    return hidden
Apply the Transformer output layer.

    Args:
      hidden: <float32>[batch_size, len, dim].
      x: <float32>[batch_size, len, dim].

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