src/transformers/models/marian/modeling_flax_marian.py/FlaxMarianEncoder/setup
class FlaxMarianEncoder:    def setup(self):
        self.dropout_layer = nn.Dropout(rate=self.config.dropout)

        embed_dim = self.config.d_model
        self.max_source_positions = self.config.max_position_embeddings
        self.embed_scale = math.sqrt(embed_dim) 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 = FlaxMarianEncoderLayerCollection(self.config, self.dtype)