Generating with Flax fails when using Causal Language models
Author: SamKG
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.21.1
- Platform: Linux-4.18.0-372.19.1.el8_6.x86_64-x86_64-with-glibc2.28
- Python version: 3.10.4
- Huggingface_hub version: 0.8.1
- PyTorch version (GPU?): 1.12.1+cu102 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): 0.6.0 (gpu)
- Jax version: 0.3.17
- JaxLib version: 0.3.15
- Using GPU in script?: Yes (Nvidia A100)
- Using distributed or parallel set-up in script?: No


### Who can help?

@patrickvonplaten 
@Narsil 
@patil-suraj 

### Information

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

### Tasks

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

### Reproduction

Run the following code snippet:

```python
from jax import numpy as jnp
import transformers

model = transformers.FlaxAutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B")
tokenizer = transformers.AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")

entence = "Paris is one of the densest populated areas in Europe."
input_ids = tokenizer(sentence, return_tensors="jax")["input_ids"]

model.generate(input_ids)
```

### Expected behavior

Expected behavior is that the model generates completions for the given input id.

Observed behavior is that the following error is thrown:
```
File ~/.conda/envs/lm-extraction/lib/python3.10/site-packages/jax/_src/lax/lax.py:4577, in _check_same_dtypes(name, ignore_fp_precision, *ttypes)
   4575   equiv = _JNP_FUNCTION_EQUIVALENTS[name]
   4576   msg += f" (Tip: jnp.{equiv} is a similar function that does automatic type promotion on inputs)."
-> 4577 raise TypeError(msg.format(name, ", ".join(map(str, types))))

TypeError: lax.dynamic_update_slice requires arguments to have the same dtypes, got int32, float32.
```

This seems to be a type mismatch error

Comments:
Hey, I just ran into the same issue.
For me it does not occur when I explicitly specify the `pad_token_id` as in this example:
```
model.generate(
    prompt_tokenized,
    params=model.params,
    pad_token_id=50256,
)
```
I'm also using the `"EleutherAI/gpt-neo-1.3B"` model, one might need to adjust the `pad_token_id` for different models / tokenizers.


You can also use the workaround as in the test here: https://github.com/huggingface/transformers/blob/a541d97477a8901e37e5f850f2cd707ffc82445b/tests/generation/test_generation_flax_utils.py#L83-L85

-> just set `model.config.pad_token_id = model.config.eos_token_id` before generating.

Hey @SamKG!

As @maxidl has kindly pointed out, the `pad_token_id` needs to be specified for autoregressive generation.

You can set this to `tokenizer.pad_token_id` **if** the tokenizer has a `pad_token_id` defined:
```python
if tokenizer.pad_token_id is not None:
    model.config.pad_token_id = tokenizer.pad_token_id
```

Otherwise, setting it to the `eos_token_id` is possible:
```python
model.config.pad_token_id = model.config.eos_token_id
```

The `pad_token_id` should always either be passed to the `.generate()` method or specified in the `config`. 

@patrickvonplaten IMO it's worth a warning message when the `pad_token_id` is omitted! Would prevent hidden errors such as this one.

Also cc'ing @patil-suraj who I believe has come across this issue before with GPT-J models

@sanchit-gandhi added to the `generate` to do list 👍 

Thanks @gante 

#21009 Fixes it -- Flax now assumes the value of `pad_token_id` when it is `None` and `eos_token_id` is not `None`, like TF and PT do. This should also be the case in the examples above.

@SamKG @maxidl I'm closing this issue as it seems to be solved, but feel free to reopen it with further queries :)