Fix incomplete outputs of FlaxBert (#18772)

* Fix incomplete FlaxBert outputs

* fix big_bird electra roberta

src/transformers/models/bert/modeling_flax_bert.py

@@ -612,7 +612,7 @@ class FlaxBertLayerCollection(nn.Module):
         if output_hidden_states:
             all_hidden_states += (hidden_states,)
 
-        outputs = (hidden_states,)
+        outputs = (hidden_states, all_hidden_states, all_attentions, all_cross_attentions)
 
         if not return_dict:
             return tuple(v for v in outputs if v is not None)

src/transformers/models/big_bird/modeling_flax_big_bird.py

@@ -1438,7 +1438,7 @@ class FlaxBigBirdLayerCollection(nn.Module):
         if output_hidden_states:
             all_hidden_states += (hidden_states,)
 
-        outputs = (hidden_states,)
+        outputs = (hidden_states, all_hidden_states, all_attentions, all_cross_attentions)
 
         if not return_dict:
             return tuple(v for v in outputs if v is not None)

src/transformers/models/electra/modeling_flax_electra.py

@@ -590,7 +590,7 @@ class FlaxElectraLayerCollection(nn.Module):
         if output_hidden_states:
             all_hidden_states += (hidden_states,)
 
-        outputs = (hidden_states,)
+        outputs = (hidden_states, all_hidden_states, all_attentions, all_cross_attentions)
 
         if not return_dict:
             return tuple(v for v in outputs if v is not None)

src/transformers/models/roberta/modeling_flax_roberta.py

@@ -580,7 +580,7 @@ class FlaxRobertaLayerCollection(nn.Module):
         if output_hidden_states:
             all_hidden_states += (hidden_states,)
 
-        outputs = (hidden_states,)
+        outputs = (hidden_states, all_hidden_states, all_attentions, all_cross_attentions)
 
         if not return_dict:
             return tuple(v for v in outputs if v is not None)