Add support for beam search's num_return_sequencs flag in flax 
Fixes part of https://github.com/huggingface/transformers/issues/22696
CC @gianlucadetommaso and @gante <|||||>_The documentation is not available anymore as the PR was closed or merged._<|||||>@mayankagarwals to make our CI green, you will likely need to rebase with `main`, then run `make fixup`, then commit. 

Also, tag the PR as ready when you're ready for a final check from a core maintainer :)<|||||>Hey @gante ,

Yes! Have done those. Wanted to get your views on it before cleaning up the PR. The CI is green now. 

I couldn't find any specific test for this so didn't add but the following script serves as a decent test for functional purposes
```
from transformers import FlaxGPT2LMHeadModel, GPT2Tokenizer

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = FlaxGPT2LMHeadModel.from_pretrained("gpt2", pad_token_id=tokenizer.eos_token_id)

input_ids = tokenizer.encode('I enjoy walking with my cute dog', return_tensors='jax')

beam_output = model.generate(
    input_ids,
    max_length=50,
    num_beams=5,
    no_repeat_ngram_size=2,
    num_return_sequences=2,
    early_stopping=True
)
print("All generated hypotheses:\n")
for sequence in beam_output.sequences.tolist():
    print(tokenizer.decode(sequence, skip_special_tokens=True))
    print("-------")
    
```
    
    
   Output before the change: 
   
```
All generated hypotheses:

I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with him again.

I'm not sure if I'll ever be able to walk with him again. I'm not sure if I'll
-------
```

Output after change: 
```
All generated hypotheses:

I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with him again.

I'm not sure if I'll ever be able to walk with him again. I'm not sure if I'll
-------
I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with him again.

I'm not sure if I'll ever be able to walk with him again.

I'm not sure if
-------
```<|||||>> But augment it to return N beams and verify we get the sequences/scores for these beams

Sure, have added the test. Thanks for the reference, made it a 5-minute work :p  @sanchit-gandhi @gante <|||||>ready for a final check, tagging a core maintainer