src/transformers/models/gptj/modeling_flax_gptj.py/FlaxGPTJBlock/setup
class FlaxGPTJBlock:    def setup(self):
        hidden_size = self.config.hidden_size
        inner_dim = self.config.n_inner if self.config.n_inner is not None else 4 * hidden_size

        self.ln_1 = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)
        self.attn = FlaxGPTJAttention(self.config, dtype=self.dtype)

        self.mlp = FlaxGPTJMLP(self.config, inner_dim, dtype=self.dtype)