src/transformers/models/roberta_prelayernorm/modeling_flax_roberta_prelayernorm.py/FlaxRobertaPreLayerNormModule/setup
class FlaxRobertaPreLayerNormModule:    def setup(self):
        self.embeddings = FlaxRobertaPreLayerNormEmbeddings(self.config, dtype=self.dtype)
        self.encoder = FlaxRobertaPreLayerNormEncoder(
            self.config,
            dtype=self.dtype,
            gradient_checkpointing=self.gradient_checkpointing,
        )
        self.LayerNorm = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype)
        self.pooler = FlaxRobertaPreLayerNormPooler(self.config, dtype=self.dtype)