src/transformers/models/bert/modeling_flax_bert.py/FlaxBertPooler/__call__
class FlaxBertPooler:    def __call__(self, hidden_states):
        cls_hidden_state = hidden_states[:, 0]
        cls_hidden_state = self.dense(cls_hidden_state)
        return nn.tanh(cls_hidden_state)