src/transformers/models/gpt2/modeling_flax_gpt2.py/FlaxGPT2MLP/setup
class FlaxGPT2MLP:    def setup(self):
        embed_dim = self.config.hidden_size
        self.c_fc = FlaxConv1D(self.intermediate_size, dtype=self.dtype)
        self.c_proj = FlaxConv1D(embed_dim, dtype=self.dtype)
        self.act = ACT2FN[self.config.activation_function]
        self.dropout = nn.Dropout(rate=self.config.resid_pdrop)