src/transformers/models/roberta_prelayernorm/modeling_flax_roberta_prelayernorm.py/FlaxRobertaPreLayerNormLayerCollection/setup
class FlaxRobertaPreLayerNormLayerCollection:    def setup(self):
        if self.gradient_checkpointing:
            FlaxRobertaPreLayerNormCheckpointLayer = remat(FlaxRobertaPreLayerNormLayer, static_argnums=(5, 6, 7))
            self.layers = [
                FlaxRobertaPreLayerNormCheckpointLayer(self.config, name=str(i), dtype=self.dtype)
                for i in range(self.config.num_hidden_layers)
            ]
        else:
            self.layers = [
                FlaxRobertaPreLayerNormLayer(self.config, name=str(i), dtype=self.dtype)
                for i in range(self.config.num_hidden_layers)
            ]