src/transformers/models/bart/modeling_flax_bart.py/FlaxBartDecoderWrapper/setup
class FlaxBartDecoderWrapper:    def setup(self):
        embed_dim = self.config.d_model
        embed_tokens = nn.Embed(
            self.config.vocab_size,
            embed_dim,
            embedding_init=jax.nn.initializers.normal(self.config.init_std),
            dtype=self.dtype,
        )
        self.decoder = FlaxBartDecoder(config=self.config, embed_tokens=embed_tokens, dtype=self.dtype)