src/transformers/models/bert/modeling_flax_bert.py/FlaxBertModule/setup
class FlaxBertModule:    def setup(self):
        self.embeddings = FlaxBertEmbeddings(self.config, dtype=self.dtype)
        self.encoder = FlaxBertEncoder(self.config, dtype=self.dtype)
        self.pooler = FlaxBertPooler(self.config, dtype=self.dtype)