`no_repeat_ngram_size` has no effect for Flax model
### System Info

transformes = ^4.27.4, macOS, python = ^3.9.6

### Who can help?

@sanchit-gandhi @ArthurZucker @younesbelkada 

### Information

- [ ] The official example scripts
- [X] 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

I have been walking through the generation example in [here](https://huggingface.co/blog/how-to-generate), but I am trying to use `FlaxGPT2LMHeadModel` instead of `GPT2LMHeadModel`. 

Everything works up to when `no_repeat_ngram_size` is introduced. In the example, setting `no_repeat_ngram_size=2` changes the generated sentence from 

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

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

I've been thinking about this for a while now, and I think it's time for me to take a break
```

However, when using `FlaxGPT2LMHeadModel` instead of `GPT2LMHeadModel`, the generated sentence with `no_repeat_ngram_size=2` remains exactly the same as the first message.

Here is a reproducing example.

```
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, 
    early_stopping=True
)

print(tokenizer.decode(beam_output.sequences.tolist()[0], skip_special_tokens=True))
# 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
```

Perhaps there is a bug in the interplay of `FlaxGPT2LMHeadModel` and `no_repeat_ngram_size`?

### Expected behavior

Everything works fine when `GPT2LMHeadModel` is used instead. Here an example.

```
from transformers import GPT2LMHeadModel, GPT2Tokenizer

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

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

beam_output = model.generate(
    input_ids, 
    max_length=50, 
    num_beams=5, 
    no_repeat_ngram_size=2, 
    early_stopping=True
)

print(tokenizer.decode(beam_output[0], skip_special_tokens=True))
# I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with him again.
# I've been thinking about this for a while now, and I think it's time for me to take a break
```
cc @sanchit-gandhi <|||||>Hey @gianlucadetommaso 

You are right. `no_repeat_ngram` is one of the various logit processors used during generation. While I can see it's defined in tensorflow : https://github.com/huggingface/transformers/blob/151425ddb29d4ad1a121e8cce62000a2ac52d3ba/src/transformers/generation/tf_utils.py#L1451 

I don't think it's defined in flax: 
https://github.com/huggingface/transformers/blob/151425ddb29d4ad1a121e8cce62000a2ac52d3ba/src/transformers/generation/flax_utils.py#L488

@sanchit-gandhi 

Let me know if you have your hands full. 
Been meaning to get into flax for a while, can take this up. Shouldn't be very problematic. I'll just need to see how to implement the same processor in flax,  might need a little guidance

<|||||>@mayankagarwals @sanchit-gandhi It seems also `num_return_sequences` does not work for the Flax model, and indeed it seems missing from [transformers/src/transformers/generation/flax_utils.py](https://github.com/huggingface/transformers/blob/151425ddb29d4ad1a121e8cce62000a2ac52d3ba/src/transformers/generation/flax_utils.py#L488). 

Worth doing a more general comparison of functionalities, I guess. 😄 <|||||>Hey, I noticed this issue randomly, so just dropping in to say it looks like #18707.

The implementation was attempted in #18769 (cc @gante) and dropped because it was memory heavy, but I think it should be doable with a XLA `while` loop without allocating huge tensors. I implemented similar logic in Elixir recently (with the same shape restrictions, since we also use XLA), so perhaps [this](https://github.com/elixir-nx/bumblebee/blob/6ae97b2ce2e99a863a658f0730334c0a4984fc3d/lib/bumblebee/text/generation.ex#L745-L779) helps. The code is fairly high-level, but if anything is not clear let me know :)<|||||>@jonatanklosko nice, I had given up on it after spending a few hours on it! I'll keep tabs on your implementation, in case no one picks it up in the near future :)

@gianlucadetommaso @mayankagarwals feel free to pick up the task of adding missing logits processors to FLAX! In general, a close copy of TF's implementations will work on FLAX, since they both rely on XLA and have similar syntax.
<|||||>@gante Got it. I'll take some time out and look into this. Quite interesting!<|||||>@gante 

While I was able to understand and reproduce your solution (kudos on the clean code), I had a question. 
The following code works as expected
```import tensorflow as tf
batch_size = 5
ngram_size = 4
vocab_size = 52
seq_len = 50


input_ids = tf.convert_to_tensor([[   40,  28,  35,   36,   37, 14,  15, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51], [   40,  28,  35,   36,   37, 14,  15, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51],[   40,  28,  35,   36,   37, 14,  15, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51], [   40,  28,  35,   36,   37, 14,  15, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51],[   40,  28,  35,   36,   37, 14,  15, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51, 51, 51, 51, 51,
       51, 51, 51, 51, 51]])
transition_tensor = tf.zeros((batch_size, ngram_size - 1, vocab_size, vocab_size), dtype=tf.bool)

# if `input_ids` is padded this will do some useless computations, but that is fine (avoids XLA recompilation)
for i in range(seq_len - (ngram_size - 1)):
    ngrams = input_ids[:, i : i + ngram_size]

    # creates the indexing for the batch and the n-th member of the ngram
    batch_indexing, ngram_indexing = tf.meshgrid(tf.range(ngrams.shape[0]), tf.range(ngrams.shape[1] - 1))
    batch_indexing = tf.reshape(tf.transpose(batch_indexing), (-1,))
    ngram_indexing = tf.reshape(tf.transpose(ngram_indexing), (-1,))

    # creates the indexing for the current -> next token p  airs
    curr_tokens = ngrams[:, :-1]
    next_tokens = ngrams[:, 1:]
    current_token_indexing = tf.reshape(curr_tokens, (-1,))
    next_token_indexing = tf.reshape(next_tokens, (-1,))

    # scatters the observed ngrams into the transition tensor
    update_indices = tf.stack(
        (batch_indexing, ngram_indexing, current_token_indexing, next_token_indexing), axis=1
    )

    transition_tensor = tf.tensor_scatter_nd_update(
        tensor=transition_tensor,
        indices=update_indices,
        updates=tf.ones(update_indices.shape[0], dtype=tf.bool),
    )

print(transition_tensor)
```

But when done with higher dimensions, gives core dump on CPU. This is just eager evaluation without any xla. Any idea why this might be? Does tensorflow give core dump when it can't do large operations on CPU? 
```# batch_size = 5
# ngram_size = 4
# vocab_size = 50257
# seq_len = 50
# 
# 
# input_ids = tf.convert_to_tensor([[   40,  2883,  6155,   351,   616, 13779,  3290, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256], [   40,  2883,  6155,   351,   616, 13779,  3290, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256],[   40,  2883,  6155,   351,   616, 13779,  3290, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256], [   40,  2883,  6155,   351,   616, 13779,  3290, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256],[   40,  2883,  6155,   351,   616, 13779,  3290, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,
#        50256, 50256, 50256, 50256, 50256]])
```<|||||>@mayankagarwals I believe I got stuck precisely at that point :) Have you had a chance to look at @jonatanklosko's implementation? <|||||>I just had a high-level overview, yet to dig in. I figured I'd first understand why this implementation wasn't working. I'll get to it soon

That is very odd behavior though. If there was an issue with resources, the code should have failed while allocating the transition tensor 
`transition_tensor = tf.zeros((batch_size, ngram_size - 1, vocab_size, vocab_size), dtype=tf.bool)
`

But instead, it fails while performing the scatter and update op. 
To debug further I broke down the scatter and update it into 15 different operations in a loop. It's failing for some and not failing for others. 
```
    for j in update_indices:
        print(j)
        tf.tensor_scatter_nd_update(transition_tensor, tf.expand_dims(j, axis=0), tf.constant([True], dtype=tf.bool))

```

Anyway, if by any chance you figure this out during your stint with TensorFlow. Do let me know, I'd be interested to know
<|||||>Keeping this open in case you want to continue on @mayankagarwals - seems like you were making good progress!<|||||>Thanks, @sanchit-gandhi. I will find time to work on this soon, just got caught up with other things. Will update the thread as I make progress :) <|||||>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.