src/transformers/models/t5/modeling_flax_t5.py/FlaxT5DenseReluDense/setup
class FlaxT5DenseReluDense:    def setup(self):
        wi_init_std = self.config.initializer_factor * (self.config.d_model ** -0.5)
        wo_init_std = self.config.initializer_factor * (self.config.d_ff ** -0.5)

        self.wi = nn.Dense(
            self.config.d_ff,
            use_bias=False,
            kernel_init=jax.nn.initializers.normal(wi_init_std, self.dtype),
            dtype=self.dtype,
        )
        self.wo = nn.Dense(
            self.config.d_model,
            use_bias=False,
            kernel_init=jax.nn.initializers.normal(wo_init_std, self.dtype),
            dtype=self.dtype,
        )
        self.dropout = nn.Dropout(self.config.dropout_rate)