Repository: huggingface/transformers

Issue:
inconsistency of the last element in hidden_states between PyTorch/Flax GPT2(Neo)
### Who can help
@patrickvonplaten  @patil-suraj 

## Information

The current Flax version of GPT2/GPTNeo give different results for the last element in `hidden_states` if `output_hidden_states=True`. This difference comes from the following fact:

In Flax GPT2 (and GPTNeo similarly), `all_hidden_states` is prepared in `FlaxGPT2BlockCollection` which has no layer norm layer (`ln_f`), therefore the last hidden state is added before applying layer normalization.
While in PyTorch/TF GPT2, it is prepared in `GPT2Model` or `TFGPT2MainLayer`, which contain `ln_f` layer, and the last hidden state is added after applying layer normalization.

This could be fixed by updating the outputs in `FlaxGPT2Module.__call__`, (if it's worth the change), something like

```
        hidden_states = outputs[0]
        hidden_states = self.ln_f(hidden_states)
        
        all_hidden_states = None
        if output_hidden_states:
            if not return_dict:
                all_hidden_states = outputs[1]
            else:
                all_hidden_states = outputs.hidden_states
            all_hidden_states = all_hidden_states[:-1] + (hidden_states,)

        if not return_dict:
            if all_hidden_states:
                return (hidden_states, all_hidden_states) + outputs[2:]
            else:
                return (hidden_states,) + outputs[1:]

        return FlaxBaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            hidden_states=all_hidden_states,
            attentions=outputs.attentions,
            cross_attentions=outputs.cross_attentions,
        )
``` 

### Related places in the source code

PyTroch GPT2
https://github.com/huggingface/transformers/blob/773d386041b2761204dcc67b316904d8d5b412da/src/transformers/models/gpt2/modeling_gpt2.py#L820

```
        hidden_states = self.ln_f(hidden_states)
        ...
        # Add last hidden state
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)
```

TensorFlow GPT2
https://github.com/huggingface/transformers/blob/773d386041b2761204dcc67b316904d8d5b412da/src/transformers/models/gpt2/modeling_tf_gpt2.py#L397

```
        hidden_states = self.ln_f(hidden_states)
        ...
        # Add last hidden state
        if inputs["output_hidden_states"]:
            all_hidden_states = all_hidden_states + (hidden_states,)
```

Flax GPT2
https://github.com/huggingface/transformers/blob/773d386041b2761204dcc67b316904d8d5b412da/src/transformers/models/gpt2/modeling_flax_gpt2.py#L461

```
       # In `FlaxGPT2BlockCollection` which has no `ln_f` (only exist in `FlaxGPT2Module`)
     
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

```


Discussion:
Hey @ydshieh,

That's a great catch! And we should definitely correct this. The way to go here in my opinion is to remove the second:

```python
if output_hidden_states:
    all_hidden_states += (hidden_states,)
```

of ```FlaxGPT2BlockCollection``` 

and move it to:

```FlaxGPT2Module``` 

as you've suggested I think. Would you be interested in opening a PR for this? :-)
Hi, yes, I can open a PR for this. But just to be sure

> my opinion is to remove the second

do you mean, in `FlaxGPT2BlockCollection`, we should (if specified) return the tuple containing all the hidden states EXCEPT the last one? And add the last one in `FlaxGPT2Module`.

I am OK with it - it is just slightly different from what I wrote originally (still add the last one in `FlaxGPT2BlockCollection`, but updating later).
Yes I think we should add all hidden states EXCEPT the last one. This class is never used externally without using `FlaxGPT2Module` so it's safe to do IMO. Adding it once is the better option instead of adding it and updating it later IMO.

Resolving patch:
diff --git a/src/transformers/models/gpt2/modeling_flax_gpt2.py b/src/transformers/models/gpt2/modeling_flax_gpt2.py
--- a/src/transformers/models/gpt2/modeling_flax_gpt2.py
+++ b/src/transformers/models/gpt2/modeling_flax_gpt2.py
@@ -24,7 +24,7 @@
 from jax import lax
 
 from ...file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
-from ...modeling_flax_outputs import FlaxBaseModelOutput, FlaxBaseModelOutputWithPast, FlaxCausalLMOutput
+from ...modeling_flax_outputs import FlaxBaseModelOutput, FlaxCausalLMOutput
 from ...modeling_flax_utils import ACT2FN, FlaxPreTrainedModel, append_call_sample_docstring
 from ...utils import logging
 from .configuration_gpt2 import GPT2Config
@@ -458,20 +458,10 @@ def __call__(
             if output_attentions:
                 all_attentions += (layer_outputs[1],)
 
-        if output_hidden_states:
-            all_hidden_states += (hidden_states,)
-
-        outputs = (hidden_states,)
-
-        if not return_dict:
-            return tuple(v for v in outputs if v is not None)
+        # this contains possible `None` values - `FlaxGPT2Module` will filter them out
+        outputs = (hidden_states, all_hidden_states, all_attentions)
 
-        return FlaxBaseModelOutputWithPast(
-            last_hidden_state=hidden_states,
-            past_key_values=None,
-            hidden_states=all_hidden_states,
-            attentions=all_attentions,
-        )
+        return outputs
 
 
 class FlaxGPT2Module(nn.Module):
@@ -527,13 +517,19 @@ def __call__(
         hidden_states = outputs[0]
         hidden_states = self.ln_f(hidden_states)
 
+        if output_hidden_states:
+            all_hidden_states = outputs[1] + (hidden_states,)
+            outputs = (hidden_states, all_hidden_states) + outputs[2:]
+        else:
+            outputs = (hidden_states,) + outputs[1:]
+
         if not return_dict:
-            return (hidden_states,) + outputs[1:]
+            return tuple(v for v in outputs if v is not None)
 
         return FlaxBaseModelOutput(
             last_hidden_state=hidden_states,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
+            hidden_states=outputs[1],
+            attentions=outputs[-1],
         )
 
 
diff --git a/src/transformers/models/gpt_neo/modeling_flax_gpt_neo.py b/src/transformers/models/gpt_neo/modeling_flax_gpt_neo.py
--- a/src/transformers/models/gpt_neo/modeling_flax_gpt_neo.py
+++ b/src/transformers/models/gpt_neo/modeling_flax_gpt_neo.py
@@ -25,7 +25,7 @@
 from jax import lax
 
 from ...file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
-from ...modeling_flax_outputs import FlaxBaseModelOutput, FlaxBaseModelOutputWithPast, FlaxCausalLMOutput
+from ...modeling_flax_outputs import FlaxBaseModelOutput, FlaxCausalLMOutput
 from ...modeling_flax_utils import ACT2FN, FlaxPreTrainedModel, append_call_sample_docstring
 from ...utils import logging
 from .configuration_gpt_neo import GPTNeoConfig
@@ -488,20 +488,10 @@ def __call__(
             if output_attentions:
                 all_attentions += (layer_outputs[1],)
 
-        if output_hidden_states:
-            all_hidden_states += (hidden_states,)
-
-        outputs = (hidden_states,)
-
-        if not return_dict:
-            return tuple(v for v in outputs if v is not None)
+        # this contains possible `None` values - `FlaxGPTNeoModule` will filter them out
+        outputs = (hidden_states, all_hidden_states, all_attentions)
 
-        return FlaxBaseModelOutputWithPast(
-            last_hidden_state=hidden_states,
-            past_key_values=None,
-            hidden_states=all_hidden_states,
-            attentions=all_attentions,
-        )
+        return outputs
 
 
 class FlaxGPTNeoModule(nn.Module):
@@ -557,13 +547,22 @@ def __call__(
         hidden_states = outputs[0]
         hidden_states = self.ln_f(hidden_states)
 
+        hidden_states = outputs[0]
+        hidden_states = self.ln_f(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states = outputs[1] + (hidden_states,)
+            outputs = (hidden_states, all_hidden_states) + outputs[2:]
+        else:
+            outputs = (hidden_states,) + outputs[1:]
+
         if not return_dict:
-            return (hidden_states,) + outputs[1:]
+            return tuple(v for v in outputs if v is not None)
 
         return FlaxBaseModelOutput(
             last_hidden_state=hidden_states,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
+            hidden_states=outputs[1],
+            attentions=outputs[-1],
         )
 
 


Test patch: