src/transformers/models/electra/modeling_flax_electra.py/FlaxElectraModule/setup
class FlaxElectraModule:    def setup(self):
        self.embeddings = FlaxElectraEmbeddings(self.config, dtype=self.dtype)
        if self.config.embedding_size != self.config.hidden_size:
            self.embeddings_project = nn.Dense(self.config.hidden_size)
        self.encoder = FlaxElectraEncoder(self.config, dtype=self.dtype)