src/transformers/models/t5/modeling_flax_t5.py/FlaxT5LayerFF/setup
class FlaxT5LayerFF:    def setup(self):
        if self.config.feed_forward_proj == "relu":
            self.DenseReluDense = FlaxT5DenseReluDense(self.config, dtype=self.dtype)
        elif self.config.feed_forward_proj == "gated-gelu":
            self.DenseReluDense = FlaxT5DenseGatedGeluDense(self.config, dtype=self.dtype)
        else:
            raise ValueError(
                f"{self.config.feed_forward_proj} is not supported. Choose between `relu` and `gated-gelu`"
            )

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