src/transformers/models/t5/modeling_flax_t5.py/FlaxT5Stack/setup
class FlaxT5Stack:    def setup(self):
        self.causal = self.config.causal

        if self.embed_tokens is None:
            self.embed_tokens = nn.Embed(
                self.config.vocab_size,
                self.config.d_model,
                embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),
                dtype=self.dtype,
            )

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