src/transformers/models/esm/modeling_tf_esm.py/TFEsmLMHead/__init__
class TFEsmLMHead:    def __init__(self, config, name=None):
        super().__init__(name=name)
        self.dense = Dense(
            config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
        )

        self.layer_norm = LayerNormalization(epsilon=config.layer_norm_eps, name="layer_norm")

        self.decoder = Dense(
            config.vocab_size,
            use_bias=False,
            kernel_initializer=get_initializer(config.initializer_range),
            name="decoder",
        )
        self.config = config