src/transformers/models/electra/modeling_flax_electra.py/FlaxElectraDiscriminatorPredictions/setup
class FlaxElectraDiscriminatorPredictions:    def setup(self):
        self.dense = nn.Dense(self.config.hidden_size, dtype=self.dtype)
        self.dense_prediction = nn.Dense(1, dtype=self.dtype)