FlaxT5ForConditionalGeneration: Inconsistency in Final Block Hidden State of Encoder/Decoder
Author: ztjhz
Pull request: no

### System Info

I'm implementing my own T5 model in JAX and using the `FlaxT5ForConditionalGeneration` module to evaluate the results of my work.

During my testing phase, I ran into an issue. Using the provided code, I noticed that the hidden states from block `0` to block `10` in my implementation are consistent with the corresponding hidden states in the transformer model (i.e., `output_flax['hidden_states'][0]` to `output_flax['hidden_states'][10]`). 

However, the issue arises in the final block, where the hidden state of my model doesn't match with the transformer model's corresponding hidden state (`output_flax['hidden_states'][11]`). This is strange because after I apply the RMS layer normalization on my final block hidden state to get the `final_hidden_state`, it aligns with the `final_hidden_state` of the transformer model (`output_flax['final_hidden_state']`).

According to my understanding, the encoder block is replicated 12 times without any special processing in the final block. Hence, I am unclear about what could be causing this inconsistency in the final block's hidden states for both the encoder and decoder.

In summary, here's what I've observed:
- My hidden state aligns with `output_flax['hidden_states'][0]` to `output_flax['hidden_states'][10]`.
- My hidden state doesn't match `output_flax['hidden_states'][11]` (before applying the final layer norm).
- My final hidden state (after applying the layer norm) aligns with `output_flax['final_hidden_state']`.

### Who can help?

_No response_

### Information

- [X] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Here is the Python code I used for testing:

```python
from transformers import AutoTokenizer, FlaxT5ForConditionalGeneration

tokenizer = AutoTokenizer.from_pretrained("t5-base")
model = FlaxT5ForConditionalGeneration.from_pretrained("allenai/unifiedqa-t5-base")

inputs = tokenizer(
    ["summarize: My friends are cool but they eat too many carbs."], return_tensors="np"
)
input_ids = inputs["input_ids"]

output_flax = model.encode(
    input_ids, output_hidden_states=True, return_dict=True, output_attentions=True
)
```

### Expected behavior

I expect that each block's hidden state in my implementation of the encoder/decoder should align with the corresponding block's hidden state in the transformer model.

Comments:
cc @sanchit-gandhi 

Hey @ztjhz - thanks for the descriptive issue! We are indeed missing a bit of functionality here - for layers 0 to 10 we append the output of the layer to our collection of hidden-states `all_hidden_states`:
https://github.com/huggingface/transformers/blob/17846646f230fdf5b0d6b4d31248ccb418000acb/src/transformers/models/t5/modeling_flax_t5.py#L691-L692
However, we **don't** add the `hidden_states` to our collection `all_hidden_states` after the last layer, we simply return all of the collections we have so far:
https://github.com/huggingface/transformers/blob/17846646f230fdf5b0d6b4d31248ccb418000acb/src/transformers/models/t5/modeling_flax_t5.py#L721

What we should do is append the `hidden_states` to our collection `all_hidden_states` after the final layer, as we do for FlaxBart:
https://github.com/huggingface/transformers/blob/17846646f230fdf5b0d6b4d31248ccb418000acb/src/transformers/models/bart/modeling_flax_bart.py#L496-L497

Would you like to open a PR to add these lines into FlaxT5?

cc @amyeroberts here - essentially, what's happening in Flax T5 is that when `output_hidden_states` is True, we append the hidden states for layers 1 to N-1:
https://github.com/huggingface/transformers/blob/17846646f230fdf5b0d6b4d31248ccb418000acb/src/transformers/models/t5/modeling_flax_t5.py#L691-L692

But then **not** for the final layer N. For this final layer, we only append the hidden state after it's gone through an extra layer norm:https://github.com/huggingface/transformers/blob/17846646f230fdf5b0d6b4d31248ccb418000acb/src/transformers/models/t5/modeling_flax_t5.py#L774

I suggested that we fix this to add the pre layer-norm hidden states to `all_hidden_states` for the Nth layer, but this would also require updating the **PyTorch** code to follow suit, since the bug is actually inherited from the PT code:
https://github.com/huggingface/transformers/blob/17846646f230fdf5b0d6b4d31248ccb418000acb/src/transformers/models/t5/modeling_t5.py#L1066-L1067

For the base T5 model, which has 6 encoder layers, we go from outputting 6 encoder hidden states (hidden states from layers 1 to 5, and then the post layer-norm last hidden state) to 7 hidden states (hidden states from layers 1 to 6, and then the post layer-norm last hidden state).

Since T5 is a very core model, WDYT about this update? Is it too breaking to make this fix?

@sanchit-gandhi @ztjhz Thanks for the detailed outline of the issue! 

T5 is by far one of our most downloaded models, and many other models in the library copy from it. I wouldn't be in favour of changing the output like this unless it was causing major issues and reported by PT users. 

Thanks for chiming in @amyeroberts! I agree that it's probably too core to change in the lib, but you can certainly modify the source file locally @ztjhz and add the changes so that you have correctness on your end

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Going to close this one since it's been established that we probably can't update such a core model in the lib with a surprise breaking change. Feel free to make the changes locally @ztjhz! You can copy the Flax T5 modelling file and make all the changes you require :)