def __call__(self, x):
    """Apply the two-layer MLP.

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

    Returns:
      output: <float32>[batch_size, ..., output_size].
    """
    x = nn.Dense(
        self.hidden_size,
        kernel_init=self.kernel_init,
        bias_init=self.bias_init,
        name='dense')(
            x)
    x = self.hidden_act(x)
    x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=self.deterministic)
    x = nn.Dense(
        self.output_size,
        kernel_init=self.kernel_init,
        bias_init=self.bias_init,
        name='dense2')(
            x)
    return x
Apply the two-layer MLP.

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

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