src/transformers/models/gptj/modeling_flax_gptj.py/FlaxGPTJForCausalLMModule/setup
class FlaxGPTJForCausalLMModule:    def setup(self):
        self.transformer = FlaxGPTJModule(self.config, dtype=self.dtype)
        self.lm_head = nn.Dense(
            self.config.vocab_size,
            dtype=self.dtype,
            kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
        )