src/transformers/models/mbart/modeling_flax_mbart.py/FlaxMBartForConditionalGenerationModule/setup
class FlaxMBartForConditionalGenerationModule:    def setup(self):
        self.model = FlaxMBartModule(config=self.config, dtype=self.dtype)
        self.lm_head = nn.Dense(
            self.model.shared.num_embeddings,
            use_bias=False,
            dtype=self.dtype,
            kernel_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),
        )
        self.final_logits_bias = self.param("final_logits_bias", self.bias_init, (1, self.model.shared.num_embeddings))