Flax OPT batch generation test
### System Info

```shell
- `transformers` version: 4.20.0.dev0
- Platform: macOS-12.4-arm64-arm-64bit
- Python version: 3.9.12
- Huggingface_hub version: 0.6.0
- PyTorch version (GPU?): 1.13.0.dev20220521 (False)
- Tensorflow version (GPU?): 2.9.0 (True)
- Flax version (CPU?/GPU?/TPU?): 0.4.2 (cpu)
- Jax version: 0.3.6
- JaxLib version: 0.3.5
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>
```


### Who can help?
@patil-suraj 
The `test_batch_generation` test in FLAX OPT (currently commented) fails.
This is due to improper handling of the padding. 

The fail output is the following  : 

```python
    def test_batch_generation(self):
        model_id = "facebook/opt-350m"
    
        tokenizer = GPT2Tokenizer.from_pretrained(model_id)
        model = FlaxOPTForCausalLM.from_pretrained(model_id)
    
        tokenizer.padding_side = "left"
    
        # use different length sentences to test batching
        sentences = [
            "Hello, my dog is a little",
            "Today, I",
        ]
    
        inputs = tokenizer(sentences, return_tensors="jax", padding=True)
        input_ids = inputs["input_ids"]
    
        outputs = model.generate(input_ids=input_ids, attention_mask=inputs["attention_mask"], trace=False)
    
        inputs_non_padded = tokenizer(sentences[0], return_tensors="jax").input_ids
        output_non_padded = model.generate(input_ids=inputs_non_padded)
    
        num_paddings = inputs_non_padded.shape[-1] - inputs["attention_mask"][-1].sum()
        inputs_padded = tokenizer(sentences[1], return_tensors="jax").input_ids
        output_padded = model.generate(input_ids=inputs_padded, max_length=model.config.max_length - num_paddings)
    
        batch_out_sentence = tokenizer.batch_decode(outputs[0], skip_special_tokens=True)
        non_padded_sentence = tokenizer.decode(output_non_padded[0][0], skip_special_tokens=True)
        padded_sentence = tokenizer.decode(output_padded[0][0], skip_special_tokens=True)
    
        expected_output_sentence = [
            "Hello, my dog is a little bit of a dork.\nI'm a little bit",
            "Today, I<s><s><s><s><s><s><s><s><s><s><s><s>"
            # TODO fix this test in next PR
            # "Today, I was in the middle of a conversation with a friend about the",
        ]
        print(batch_out_sentence, [non_padded_sentence, padded_sentence])
        self.assertListEqual(expected_output_sentence, batch_out_sentence)
        # TODO outputs will be similar, fix in next PR
       self.assertListEqual(batch_out_sentence, [non_padded_sentence, padded_sentence])
````

```python
["Hello, my dog is a little bit of a dork.\nI'm a little bit", 'Today, I<s><s><s><s><s><s><s><s><s><s><s><s>'] ["Hello, my dog is a little bit of a dork.\nI'm a little bit", 'Today, I was in the middle of a conversation with a friend about the']
E       AssertionError: Lists differ: ["Hel[62 chars]ay, I<s><s><s><s><s><s><s><s><s><s><s><s>'] != ["Hel[62 chars]ay, I was in the middle of a conversation with[16 chars]the']
E       
E       First differing element 1:
E       'Today, I<s><s><s><s><s><s><s><s><s><s><s><s>'
E       'Today, I was in the middle of a conversation with a friend about the'
E       
E         ["Hello, my dog is a little bit of a dork.\nI'm a little bit",
E       -  'Today, I<s><s><s><s><s><s><s><s><s><s><s><s>']
E       +  'Today, I was in the middle of a conversation with a friend about the']

tests/models/opt/test_modeling_flax_opt.py:406: AssertionError
``` 


### 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

Just un-comment the test in the main branch

### Expected behavior

```shell
["Hello, my dog is a little bit of a dork.\nI'm a little bit","Today, I was in the middle of a conversation with a friend about the"]
```

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.<|||||>@ArthurZucker I was going to start digging into this, please confirm if thats ok per your. earlier advice, it seems strange because the issue itself is closed so its a bit confusing :)<|||||>Hi @skanjila & @ArthurZucker. I had commented on this issue recently as I ran into the same problem with the PyTorch version of OPT when performing batch generation with half precision. However, later I found out it's also related to issue #17433 and solved in the PR #17437. I installed the latest version of the library from the `main` branch and I can confirm the issue was fixed.<|||||>Nice sorry I didn't know it was already fixed! Good job ☺️ thanks both 