inconsistency of the last element in hidden_states between PyTorch/Flax GPT2(Neo)
Author: ydshieh
Pull request: no

### 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,)

```

Comments:
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.

Hi,

While working on a PR for this, it seems there is another bug in `FlaxGPT2BlockCollection`.

Near the end of its call method,

```
        outputs = (hidden_states,)

        if not return_dict:
            return tuple(v for v in outputs if v is not None)
```

it should be 

```
        outputs = (hidden_states, all_hidden_states, all_attentions)
```

I think. Otherwise, we never get `all_hidden_states` / `all_attentions` in the tuple.
(FlaxBartModel has done the right way.)

I am going to include a fix for this into the same PR. Is it ok for you?

That's a great catch!

>I am going to include a fix for this into the same PR. Is it ok for you?

Yes! Would be great if you fix it in the same PR.