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

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

    Returns:
      output: <float32>[batch_size, len, intermediate_size].
    """
    hidden = nn.Dense(
        self.intermediate_size,
        kernel_init=self.kernel_init,
        bias_init=self.bias_init,
        name='dense')(
            hidden)
    hidden = self.hidden_act(hidden)
    return hidden
Apply the Transformer intermediate layer.

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

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