[FlaxBert] Add ForCausalLM
# What does this PR do?
Adds cross-attention blocks to the following module classes:

- FlaxBertModule
- FlaxRobertaModule (in part through copying FlaxBertModule)
- FlaxBigBirdModule (in part through copying FlaxBertModule)
- FlaxElectraModule (in part through copying FlaxBertModule)


Adds the following ForCausalLM model classes:

- FlaxBertForCausalLM
- FlaxRobertaForCausalLM (in part through copying FlaxBertForCausalLM)
- FlaxBigBirdForCausalLM (in part through copying FlaxBertForCausalLM)
- FlaxElectraForCausalLM (in part through copying FlaxBertForCausalLM)

Adds the following model tests:

- FlaxRobertaForCausalLM
- FlaxBigBirdForCausalLM
- FlaxElectraForCausalLM

Note: FlaxBertForCausalLM is excluded due to the name mismatch with the PyTorch equivalent BertLMHeadModel. It is implicitly tested through the FlaxRobertaForCausalLM model tests, as well as in the following encoder-decoder model tests:

- Bert-2-Bert (encoder-decoder)
- Wav2Vec2-2-Bert (speech encoder-decoder)
_The documentation is not available anymore as the PR was closed or merged._<|||||>> Looks good to me - @sanchit-gandhi could you check though which models don't pass with `1e-5` and ideally why?
> 
> Overall `4e-2` is fine for me though cc @ydshieh what do you think?

Keep `1e-5` is much better, because so far I can always find some issues when I find something higher than `1e-5` (well, sometimes it took quite some time to figure out)