Flax T5 model - code typo during AutoRegressive decoding?
Author: giganttheo
Pull request: no

### System Info

- `transformers` version: 4.33.3
- Platform: Linux-5.15.120+-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.3.3
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.0.1+cu118 (False)
- Tensorflow version (GPU?): 2.13.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.7.4 (cpu)
- Jax version: 0.4.10
- JaxLib version: 0.4.10
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help?

@sanchit-gandhi

### Information

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

### Tasks

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

### Reproduction

The discussed file is transformers/src/transformers/models/t5/modeling_flax_t5.py at line 408

https://github.com/huggingface/transformers/blob/2aef9a96011133f6b399b598fd69cfeca936eb37/src/transformers/models/t5/modeling_flax_t5.py#L408C1-L409C1


### Expected behavior

Hi,

During autoregressive decoding, keys and values are computed one token at a time and cache is used to recover the keys and values from previous calls.

In the `_concatenate_to_cache` method of the Attention module, an attention mask is computed in order for the new query to only attend to the previous key positions and not the remaining zero elements. This is what is explained in the comments in this function.

However, the new attention mask is not used afterwards, because its name is `attention_attention_mask` and not `attention_mask` which is the one being used in every other line. 

From my understanding, this is likely a typo and I am not sure how it changes the behavior of the model, if at all.

Comments:
Hey @giganttheo - this indeed looks like a bug! The attention mask should be overridden, as in Flax GPT Neo: https://github.com/huggingface/transformers/blob/2aef9a96011133f6b399b598fd69cfeca936eb37/src/transformers/models/gpt_neo/modeling_flax_gpt_neo.py#L219

Would you like to open a PR to correct this? We can check whether we get different outputs after the fix by running the slow tests: https://github.com/huggingface/transformers/blob/57f44dc4288a3521bd700405ad41e90a4687abc0/tests/models/t5/test_modeling_flax_t5.py#L761

Which you can do so with:
```
RUN_SLOW=1 pytest -sv tests/models/t5/test_modeling_flax_t5.py::FlaxT5ModelIntegrationTests
```

Thank you for your response, I ran the slow tests on google colab with a T4 GPU and the outputs give:

```python
=============================== warnings summary ===============================
../../usr/local/lib/python3.10/dist-packages/_pytest/config/__init__.py:1373
  /usr/local/lib/python3.10/dist-packages/_pytest/config/__init__.py:1373: PytestConfigWarning: Unknown config option: doctest_glob
  
    self._warn_or_fail_if_strict(f"Unknown config option: {key}\n")

tests/models/t5/test_modeling_flax_t5.py::FlaxT5ModelIntegrationTests::test_small_generation
tests/models/t5/test_modeling_flax_t5.py::FlaxT5ModelIntegrationTests::test_small_generation_bfloat16
  /content/transformers/src/transformers/generation/flax_utils.py:322: UserWarning: You have modified the pretrained model configuration to control generation. This is a deprecated strategy to control generation and will be removed soon, in a future version. Please use and modify the model generation configuration (see https://huggingface.co/docs/transformers/generation_strategies#default-text-generation-configuration )
    warnings.warn(

tests/models/t5/test_modeling_flax_t5.py::FlaxT5ModelIntegrationTests::test_summarization
  /content/transformers/src/transformers/models/t5/tokenization_t5.py:238: FutureWarning: This tokenizer was incorrectly instantiated with a model max length of 512 which will be corrected in Transformers v5.
  For now, this behavior is kept to avoid breaking backwards compatibility when padding/encoding with `truncation is True`.
  - Be aware that you SHOULD NOT rely on t5-base automatically truncating your input to 512 when padding/encoding.
  - If you want to encode/pad to sequences longer than 512 you can either instantiate this tokenizer with `model_max_length` or pass `max_length` when encoding/padding.
  - To avoid this warning, please instantiate this tokenizer with `model_max_length` set to your preferred value.
    warnings.warn(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
================== 6 passed, 4 warnings in 242.20s (0:04:02) ===================
```

I am not very familiar with the contribution process so I am not sure what the next steps are for this, I tried following the huggingface contribution guide (https://huggingface.co/transformers/v4.2.2/contributing.html ), do you have any other resources to learn the basics of contribution guides for open source projects and transformers? I might have other contributions to do in the near future :)

That looks good - all the slow tests have passed 👍 We can see this from the last line on the print out:
```python
================== 6 passed, 4 warnings in 242.20s (0:04:02) ===================
```

That means we get correct results with the changes. The next steps would be [opening a pull request (PR)](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request) with the correction. Feel free to tag me on the PR and I'll get you a review asap. Once reviewed, we can merge into `main`. I would mention on the PR that you made the fix and then subsequently ran the slow tests and they passed - this will inform the reviewer that the correct T5 behaviour is maintained