src/transformers/models/distilbert/modeling_flax_distilbert.py/FlaxDistilBertModule/setup
class FlaxDistilBertModule:    def setup(self):
        self.embeddings = FlaxEmbeddings(self.config, dtype=self.dtype)
        self.transformer = FlaxTransformerEncoder(self.config, dtype=self.dtype)