src/transformers/models/electra/modeling_flax_electra.py/FlaxElectraSequenceSummary/setup
class FlaxElectraSequenceSummary:    def setup(self):
        self.summary = identity
        if hasattr(self.config, "summary_use_proj") and self.config.summary_use_proj:
            if (
                hasattr(self.config, "summary_proj_to_labels")
                and self.config.summary_proj_to_labels
                and self.config.num_labels > 0
            ):
                num_classes = self.config.num_labels
            else:
                num_classes = self.config.hidden_size
            self.summary = nn.Dense(num_classes, dtype=self.dtype)

        activation_string = getattr(self.config, "summary_activation", None)
        self.activation = ACT2FN[activation_string] if activation_string else lambda x: x  # noqa F407

        self.first_dropout = identity
        if hasattr(self.config, "summary_first_dropout") and self.config.summary_first_dropout > 0:
            self.first_dropout = nn.Dropout(self.config.summary_first_dropout)

        self.last_dropout = identity
        if hasattr(self.config, "summary_last_dropout") and self.config.summary_last_dropout > 0:
            self.last_dropout = nn.Dropout(self.config.summary_last_dropout)