src/transformers/models/roberta/modeling_flax_roberta.py/FlaxRobertaLMHead/__call__
class FlaxRobertaLMHead:    def __call__(self, hidden_states, shared_embedding=None):
        hidden_states = self.dense(hidden_states)
        hidden_states = ACT2FN["gelu"](hidden_states)
        hidden_states = self.layer_norm(hidden_states)

        if shared_embedding is not None:
            hidden_states = self.decoder.apply({"params": {"kernel": shared_embedding.T}}, hidden_states)
        else:
            hidden_states = self.decoder(hidden_states)

        hidden_states += self.bias
        return hidden_states