Repository: huggingface/transformers (Python)
Task type: Bug Fix

Issue:
Problem caused by boolean attention mask in `pretrained_model.generate` of Flax GPT2
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.



Discussion:
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!

Solution patch:
diff --git a/src/transformers/models/gpt2/modeling_flax_gpt2.py b/src/transformers/models/gpt2/modeling_flax_gpt2.py
--- a/src/transformers/models/gpt2/modeling_flax_gpt2.py
+++ b/src/transformers/models/gpt2/modeling_flax_gpt2.py
@@ -753,7 +753,9 @@ def prepare_inputs_for_generation(self, input_ids, max_length, attention_mask: O
         extended_attention_mask = jnp.ones((batch_size, max_length), dtype="i4")
         if attention_mask is not None:
             position_ids = attention_mask.cumsum(axis=-1) - 1
-            extended_attention_mask = lax.dynamic_update_slice(extended_attention_mask, attention_mask, (0, 0))
+            extended_attention_mask = lax.dynamic_update_slice(
+                extended_attention_mask, attention_mask.astype("i4"), (0, 0)
+            )
         else:
             position_ids = jnp.broadcast_to(jnp.arange(seq_length, dtype="i4")[None, :], (batch_size, seq_length))
 


Test patch:
diff --git a/tests/models/gpt2/test_modeling_flax_gpt2.py b/tests/models/gpt2/test_modeling_flax_gpt2.py
--- a/tests/models/gpt2/test_modeling_flax_gpt2.py
+++ b/tests/models/gpt2/test_modeling_flax_gpt2.py
@@ -187,6 +187,26 @@ def check_use_cache_forward_with_attn_mask(self, model_class_name, config, input
         diff = np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])))
         self.parent.assertTrue(diff < 1e-3, msg=f"Max diff is {diff}")
 
+    def check_bool_attention_mask_in_generation(self, model_class_name, config, input_ids, attention_mask):
+        model = model_class_name(config)
+
+        output_int_att_mask = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=3,
+        )
+
+        output_bool_att_mask = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask.astype(bool),
+            max_new_tokens=3,
+        )
+
+        self.parent.assertTrue(
+            (output_bool_att_mask.sequences == output_int_att_mask.sequences).all(),
+            "Generated response differ between boolean and integer attention mask",
+        )
+
 
 @require_flax
 class FlaxGPT2ModelTest(FlaxModelTesterMixin, FlaxGenerationTesterMixin, unittest.TestCase):
@@ -208,6 +228,13 @@ def test_use_cache_forward_with_attn_mask(self):
                 model_class_name, config, input_ids, attention_mask
             )
 
+    def test_bool_attention_mask_in_generation(self):
+        for model_class_name in self.all_generative_model_classes:
+            config, input_ids, attention_mask = self.model_tester.prepare_config_and_inputs()
+            self.model_tester.check_bool_attention_mask_in_generation(
+                model_class_name, config, input_ids, attention_mask
+            )
+
     @slow
     def test_batch_generation(self):
         tokenizer = GPT2Tokenizer.from_pretrained("gpt2", pad_token="
", padding_side="left")