src/transformers/models/longt5/modeling_flax_longt5.py/FlaxLongT5LayerFF/setup
class FlaxLongT5LayerFF:    def setup(self):
        if self.config.is_gated_act:
            self.DenseReluDense = FlaxLongT5DenseGatedActDense(self.config, dtype=self.dtype)
        else:
            self.DenseReluDense = FlaxLongT5DenseActDense(self.config, dtype=self.dtype)

        self.layer_norm = FlaxLongT5LayerNorm(
            self.config.d_model, eps=self.config.layer_norm_epsilon, dtype=self.dtype
        )
        self.dropout = nn.Dropout(self.config.dropout_rate)