src/transformers/models/roformer/modeling_flax_roformer.py/FlaxRoFormerOnlyMLMHead/__call__
class FlaxRoFormerOnlyMLMHead:    def __call__(self, hidden_states, shared_embedding=None):
        hidden_states = self.predictions(hidden_states, shared_embedding=shared_embedding)
        return hidden_states