src/transformers/models/xglm/modeling_flax_xglm.py/FlaxXGLMForCausalLMModule/setup
class FlaxXGLMForCausalLMModule:    def setup(self):
        self.model = FlaxXGLMModule(self.config, self.dtype)
        self.lm_head = nn.Dense(
            self.config.vocab_size,
            use_bias=False,
            dtype=self.dtype,
            kernel_init=jax.nn.initializers.normal(self.config.init_std),
        )