src/transformers/models/mbart/modeling_flax_mbart.py/FlaxMBartForQuestionAnsweringModule/setup
class FlaxMBartForQuestionAnsweringModule:    def setup(self):
        self.model = FlaxMBartModule(config=self.config, dtype=self.dtype)
        self.qa_outputs = nn.Dense(
            self.num_labels, dtype=self.dtype, kernel_init=jax.nn.initializers.normal(self.config.init_std)
        )