src/transformers/models/gpt_neo/modeling_flax_gpt_neo.py/FlaxGPTNeoModule/setup
class FlaxGPTNeoModule:    def setup(self):
        self.embed_dim = self.config.hidden_size
        embedding_init = jax.nn.initializers.normal(stddev=self.config.initializer_range)
        self.wte = nn.Embed(
            self.config.vocab_size,
            self.embed_dim,
            embedding_init=embedding_init,
            dtype=self.dtype,
        )
        self.wpe = nn.Embed(
            self.config.max_position_embeddings,
            self.embed_dim,
            embedding_init=embedding_init,
            dtype=self.dtype,
        )
        self.dropout = nn.Dropout(rate=self.config.embed_dropout)
        self.h = FlaxGPTNeoBlockCollection(self.config, dtype=self.dtype)
        self.ln_f = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)