[FlaxBert] Add ForCausalLM (#16995)

* [FlaxBert] Add ForCausalLM

* make style

* fix output attentions

* Add RobertaForCausalLM

* remove comment

* fix fx-to-pt model loading

* remove comment

* add modeling tests

* add enc-dec model tests

* add big_bird

* add electra

* make style

* make repo-consitency

* add to docs

* remove roberta test

* quality

* amend cookiecutter

* fix attention_mask bug in flax bert model tester

* tighten pt-fx thresholds to 1e-5

* add 'copied from' statements

* amend 'copied from' statements

* amend 'copied from' statements

* quality

docs/source/en/model_doc/bert.mdx

@@ -166,6 +166,11 @@ This model was contributed by [thomwolf](https://huggingface.co/thomwolf). The o
 [[autodoc]] FlaxBertForPreTraining
     - __call__
 
+## FlaxBertForCausalLM
+
+[[autodoc]] FlaxBertForCausalLM
+    - __call__
+
 ## FlaxBertForMaskedLM
 
 [[autodoc]] FlaxBertForMaskedLM

docs/source/en/model_doc/big_bird.mdx

@@ -120,6 +120,11 @@ This model was contributed by [vasudevgupta](https://huggingface.co/vasudevgupta
 [[autodoc]] FlaxBigBirdForPreTraining
     - __call__
 
+## FlaxBigBirdForCausalLM
+
+[[autodoc]] FlaxBigBirdForCausalLM
+    - __call__
+
 ## FlaxBigBirdForMaskedLM
 
 [[autodoc]] FlaxBigBirdForMaskedLM

docs/source/en/model_doc/electra.mdx

@@ -158,6 +158,11 @@ This model was contributed by [lysandre](https://huggingface.co/lysandre). The o
 [[autodoc]] FlaxElectraForPreTraining
     - __call__
 
+## FlaxElectraForCausalLM
+
+[[autodoc]] FlaxElectraForCausalLM
+    - __call__
+
 ## FlaxElectraForMaskedLM
 
 [[autodoc]] FlaxElectraForMaskedLM

docs/source/en/model_doc/roberta.mdx

@@ -136,6 +136,11 @@ This model was contributed by [julien-c](https://huggingface.co/julien-c). The o
 [[autodoc]] FlaxRobertaModel
     - __call__
 
+## FlaxRobertaForCausalLM
+
+[[autodoc]] FlaxRobertaForCausalLM
+    - __call__
+
 ## FlaxRobertaForMaskedLM
 
 [[autodoc]] FlaxRobertaForMaskedLM