Change Flax GPT2 with cross-attn outputs to be the same as PyTorch's version

vit_gpt2/modeling_flax_gpt2.py

@@ -593,28 +593,21 @@ class FlaxGPT2BlockCollection(nn.Module):
         if output_hidden_states:
             all_hidden_states += (hidden_states,)
 
+        # In Flax, `past_key_values` is not contained in modules' outputs.
         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)
 
-        if encoder_hidden_states is None:
-            # only self_attn
-            return FlaxBaseModelOutputWithPast(
-                last_hidden_state=hidden_states,
-                past_key_values=None,
-                hidden_states=all_hidden_states,
-                attentions=all_attentions,
-            )
-        else:
-            # with cross_attn
-            return FlaxBaseModelOutputWithPastAndCrossAttentions(
-                last_hidden_state=hidden_states,
-                past_key_values=None,
-                hidden_states=all_hidden_states,
-                attentions=all_attentions,
-                cross_attentions=all_cross_attentions,
-            )
+        # with cross_attn
+        return FlaxBaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=None,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+            cross_attentions=all_cross_attentions,
+        )
+
 
 class FlaxGPT2Module(nn.Module):
     config: GPT2Config
@@ -676,19 +669,13 @@ class FlaxGPT2Module(nn.Module):
         if not return_dict:
             return (hidden_states,) + outputs[1:]
 
-        if encoder_hidden_states is None:
-            return FlaxBaseModelOutput(
-                last_hidden_state=hidden_states,
-                hidden_states=outputs.hidden_states,
-                attentions=outputs.attentions,
-            )
-        else:
-            return FlaxBaseModelOutputWithPastAndCrossAttentions(
-                last_hidden_state=hidden_states,
-                hidden_states=outputs.hidden_states,
-                attentions=outputs.attentions,
-                cross_attentions=outputs.cross_attentions,
-            )
+        return FlaxBaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            cross_attentions=outputs.cross_attentions,
+        )
+
 
 @add_start_docstrings(
     "The bare GPT2 Model transformer outputting raw hidden-states without any specific head on top.",
@@ -753,16 +740,13 @@ class FlaxGPT2LMHeadModule(nn.Module):
         if not return_dict:
             return (lm_logits,) + outputs[1:]
 
-        if encoder_hidden_states is None:
-            return FlaxCausalLMOutput(logits=lm_logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)
-        else:
-            return FlaxCausalLMOutputWithCrossAttentions(
-                logits=lm_logits,
-                past_key_values=None,
-                hidden_states=outputs.hidden_states,
-                attentions=outputs.attentions,
-                cross_attentions=outputs.cross_attentions
-            )
+        return FlaxCausalLMOutputWithCrossAttentions(
+            logits=lm_logits,
+            past_key_values=None,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            cross_attentions=outputs.cross_attentions
+        )
 
 @add_start_docstrings(
     """