src/transformers/models/bert/modeling_flax_bert.py/FlaxBertPooler/setup
class FlaxBertPooler:    def setup(self):
        self.dense = nn.Dense(
            self.config.hidden_size,
            kernel_init=jax.nn.initializers.normal(self.config.initializer_range),
            dtype=self.dtype,
        )