src/transformers/models/pegasus/modeling_flax_pegasus.py/FlaxPegasusDecoder/setup
class FlaxPegasusDecoder:    def setup(self):
        self.dropout_layer = nn.Dropout(rate=self.config.dropout)

        embed_dim = self.config.d_model
        self.padding_idx = self.config.pad_token_id
        self.max_target_positions = self.config.max_position_embeddings
        self.embed_scale = math.sqrt(self.config.d_model) if self.config.scale_embedding else 1.0

        if self.embed_tokens is None:
            self.embed_tokens = nn.Embed(
                self.config.vocab_size,
                embed_dim,
                embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),
                dtype=self.dtype,
            )

        self.embed_positions = create_sinusoidal_positions(
            self.config.max_position_embeddings, embed_dim, dtype=self.dtype
        )

        self.layers = FlaxPegasusDecoderLayerCollection(self.config, self.dtype)
        self.layer_norm = nn.LayerNorm(dtype=self.dtype)