src/transformers/models/roberta_prelayernorm/modeling_flax_roberta_prelayernorm.py/FlaxRobertaPreLayerNormAttention/setup
class FlaxRobertaPreLayerNormAttention:    def setup(self):
        self.self = FlaxRobertaPreLayerNormSelfAttention(self.config, causal=self.causal, dtype=self.dtype)
        self.output = FlaxRobertaPreLayerNormSelfOutput(self.config, dtype=self.dtype)
        self.LayerNorm = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype)