src/transformers/models/whisper/modeling_flax_whisper.py/FlaxWhisperDecoder/setup
class FlaxWhisperDecoder:    def setup(self) -> None:
        self.embed_tokens = nn.Embed(self.config.vocab_size, self.config.d_model, dtype=self.dtype)
        self.embed_positions = nn.Embed(self.config.max_target_positions, self.config.d_model, dtype=self.dtype)

        self.layers = FlaxWhisperDecoderLayerCollection(self.config, dtype=self.dtype)

        self.dropout_layer = nn.Dropout(rate=self.config.dropout)

        self.layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-5)