Problem caused by boolean attention mask in `pretrained_model.generate` of Flax GPT2
Author: liutianlin0121
Pull request: no

Hi! 

I notice that the usage of a boolean attention mask in `pretrained_model.generate` of Flax GPT2 can cause an error. Here is a short, self-contained code block to showcase the problem; I also prepared a [colab notebook here](https://colab.research.google.com/drive/1fIfOr0AFfWlAho1dwuk8zqxKxlKmzd7i?usp=sharing):

``` python
import transformers
import jax
import jax.numpy as jnp

tokenizer = transformers.AutoTokenizer.from_pretrained(
  "gpt2", padding_side="right")
tokenizer.pad_token = tokenizer.eos_token
query = jnp.array([
    [tokenizer.pad_token_id, tokenizer.pad_token_id, 23073],
])

response_length = 4
# temperature = 0.7
pretrained_model = transformers.FlaxAutoModelForCausalLM.from_pretrained("gpt2")
generation_config = transformers.GenerationConfig(
    max_new_tokens=response_length,
    min_new_tokens=response_length,
    do_sample=True,
)
generation_config.pad_token_id = tokenizer.pad_token_id

context_length = query.shape[1]
attention_mask = query != tokenizer.pad_token_id

input_ids = query.clone()

# set padding tokens to 0
input_ids = jnp.where(attention_mask, input_ids, 0)
output = pretrained_model.generate(
    input_ids=input_ids,
    attention_mask=attention_mask,
    generation_config=generation_config,
)

# TypeError: lax.dynamic_update_slice requires arguments to have the same dtypes, got int32, bool.
```



The type error occurs because the `attention_mask` in our example above is a boolean array. But the `extended_attention_mask` used in [here](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_flax_gpt2.py#L753) internally for response generation has an integer type. This leads to an error in the `lax.dynamic_update_slice` [line here](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_flax_gpt2.py#L756), as it can't handle inputs with different data types (integer and boolean).

I think this can be a bug, because boolean attention mask should be permitted.
To fix it, one can simply update [this line](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_flax_gpt2.py#L756) in `transformers.models.gpt2.modelling_flax_gpt2.py`, which currently reads
`extended_attention_mask = lax.dynamic_update_slice(extended_attention_mask, attention_mask, (0, 0))`
into the following new line:
`extended_attention_mask = lax.dynamic_update_slice(extended_attention_mask, attention_mask.astype("i4"), (0, 0))`
This will correct the mismatch in dtypes. 

Happy to submit a PR for that!

### Who can help?

@sanchit-gandhi, @gante

### Information

- [ ] The official example scripts
- [X] 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

Here is a short, self-contained code block to showcase the problem; I also prepared a [colab notebook here](https://colab.research.google.com/drive/1fIfOr0AFfWlAho1dwuk8zqxKxlKmzd7i?usp=sharing):

``` python
import torch
import transformers
import jax
import jax.numpy as jnp

tokenizer = transformers.AutoTokenizer.from_pretrained(
  "gpt2", padding_side="right")
tokenizer.pad_token = tokenizer.eos_token
query = jnp.array([
    [tokenizer.pad_token_id, tokenizer.pad_token_id, 23073],
])

response_length = 4
# temperature = 0.7
pretrained_model = transformers.FlaxAutoModelForCausalLM.from_pretrained("gpt2")
generation_config = transformers.GenerationConfig(
    max_new_tokens=response_length,
    min_new_tokens=response_length,
    do_sample=True,
)
generation_config.pad_token_id = tokenizer.pad_token_id

context_length = query.shape[1]
attention_mask = query != tokenizer.pad_token_id

input_ids = query.clone()

# set padding tokens to 0
input_ids = jnp.where(attention_mask, input_ids, 0)
output = pretrained_model.generate(
    input_ids=input_ids,
    attention_mask=attention_mask,
    generation_config=generation_config,
)
# TypeError: lax.dynamic_update_slice requires arguments to have the same dtypes, got int32, bool.
```

### Expected behavior

I expected to execute the line `output = pretrained_model.generate(
    input_ids=input_ids,
    attention_mask=attention_mask,
    generation_config=generation_config,
)` in the above example, when `attention_mask` is a boolean mask.


Comments:
cc @sanchit-gandhi 

Hey @liutianlin0121! Thanks for the comprehensive issue description! That's a good spot - we actually covert the `attention_mask` to `"i4"` dtype under-the-hood when we call the Flax module:
https://github.com/huggingface/transformers/blob/450a181d8b963b4e896be4aac701815aa554a6bb/src/transformers/models/gpt2/modeling_flax_gpt2.py#L510

But this happens **after** the `prepare_inputs_for_generation` method. So at the point you've mentioned, we could have multiple dtypes for the attention mask (bool or int)

Given we automatically convert the attention mask to `"i4"` when we call the Flax module, I think it's safe to assume we can also do so in the `prepare_inputs_for_generation` method. This won't be surprising for the user - there's no change to behaviour here since ultimately the attention mask will be `"i4"` anyway

Feel free to open a PR to make this change and I can get you a quick approval!

Thank you! @ArthurZucker @sanchit-gandhi  I've submitted a PR.