Flax LMHeadModel for common models like Bert and Albert
### Feature request

A LM Head Model for common models such as Bert and Albert is available in both PyTorch and TensorFlow, but it appears to be missing in Flax. 

### Motivation

We are developing a library in JAX and Flax for uncertainty quantification, and we rely on Hugging Face transformers written in Flax.

### Your contribution

I would love to contribute. However, unfortunately I might not have much time to submit a PR myself.
Alright, I found out that LMHeadModel is an old naming choice that was kept not to break compatibility. In Flax, I should be able to do what I want with `FlaxBertForCausalLM`.