[Flax BERT/Roberta] few small fixes
# What does this PR do?

This PR fixes a few docs in Flax models.
nit question: Any reason why we are accepting `bias_init` as a parameter while in all other modules we just use directly in `setup`
https://github.com/huggingface/transformers/blob/a5d2967bd8a5ed2456c593fa9eb5d9c0d726ae7a/src/transformers/models/bert/modeling_flax_bert.py#L480-L483