Fix flax gpt2 hidden states
# What does this PR do?

Fixes #13102 

## Before submitting

- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests), Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? 
    [ inconsistency of the last element in hidden_states between PyTorch/Flax GPT2(Neo) #13102 ](https://github.com/huggingface/transformers/issues/13102#issuecomment-897687182)

## Who can review?

@patrickvonplaten @patil-suraj 
Hi, @patil-suraj , thank you for the suggestions. There is however some issues regarding

```
        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:]
```
because this will change `outputs` to a `tuple` even if it is previously a `FlaxBaseModelOutput`, and this causes problem to the end (i.e. if `return_dict=True`)

```
        return FlaxBaseModelOutput(
            last_hidden_state=hidden_states,
            hidden_states=outputs.all_hidden_states,
            attentions=outputs.attentions,
        )
```

Do you have a good solution to address this while keep your suggestions?<|||||>Ahh, yeah, you're right!

I wanted to avoid multiple if/else conds, but seems we will need to add one either way.

I could see two options:

- we have already stored `all_hidden_states`, we could store the `all_attentions` using 

```python3
all_attentions = outputs[-1] if output_attentions else None
```

and then use that in the output class

- another option is, `FlaxGPT2BlockCollection` is only used internally, so we could also just always return `outputs` (including `None` values) as a `tuple`. So in the `FlaxGPT2Module`, we could do 

```python3
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 tuple(v for v in outputs if v is not None)

return FlaxBaseModelOutput(
    last_hidden_state=hidden_states,
    hidden_states=all_hidden_states,
    attentions=outputs[-1],
)
```<|||||>@patil-suraj  I went for option 2, with a slightly change `hidden_states=all_hidden_states,` -> `hidden_states=outputs[1],` (all_hidden_states not always define).