src/transformers/models/albert/modeling_flax_albert.py/FlaxAlbertModule/setup
class FlaxAlbertModule:    def setup(self):
        self.embeddings = FlaxAlbertEmbeddings(self.config, dtype=self.dtype)
        self.encoder = FlaxAlbertEncoder(self.config, dtype=self.dtype)
        if self.add_pooling_layer:
            self.pooler = nn.Dense(
                self.config.hidden_size,
                kernel_init=jax.nn.initializers.normal(self.config.initializer_range, self.dtype),
                dtype=self.dtype,
                name="pooler",
            )
            self.pooler_activation = nn.tanh
        else:
            self.pooler = None
            self.pooler_activation = None