Make TFNoRepeatNGramLogitsProcessor XLA compatible 
Author: mizhazha
Pull request: no

### Feature request

make TFNoRepeatNGramLogitsProcessor XLA compatible 

### Motivation

We would like to use the BART model for summarization tasks. We are excited about the recent announcement of fast text generation with tensorflow; however, we realized that the [TFNoRepeatNGramLogitsProcessor](https://github.com/huggingface/transformers/blob/fd9aa82b07d9b844a21f18f1622de5ca104f25bd/src/transformers/generation_tf_logits_process.py#L385) is not XLA compatible therefore we cannot compile the generate method into a TF graph yet, which makes it hard to meet our throughput requirement. 

### Your contribution

@gante Since you mentioned challenges in making this compatible with XLA, do you have an estimate on how long it will take for implementation? I can chat with the team to see how we would like to proceed, as this is a hard blocker for us. 

If we are open to not use the TFNoRepeatNGramLogitsProcessor, could you suggest some hack to avoid compiling this part? 

Comments:
Hi @mizhazha 👋 

Yeah, it's unfortunate that we weren't able to implement it with XLA. Sadly, using XLA is very binary -- either the entire function (`.generate()`) and all its inner calls can be compiled, or it is not XLA-compatible. If it is not a hard requirement for your application, the easiest solution is not to use it for now :D (i.e. don't set the `no_repeat_ngram_size` argument)

I've spent some time a few months ago trying to make it XLA-compatible, and I've had a few ideas to try to reimplement the processor with XLA capabilities. I will give it a go this week, maybe it can be solved quickly.

If I give no reply until the following Monday, feel free to ping me to brainstorm possible alternatives.

Hi @mizhazha 👋 

I have attempted to explore alternatives in [this PR](https://github.com/huggingface/transformers/pull/18769). It seems like the simple alternatives are compute and memory heavy, so not worth merging (running without XLA is faster). I still think that efficient data structures might enable an efficient implementation, but we don't have the resources to explore those alternatives.

If your team still wants to use XLA and `no_repeat_ngram_size`, a viable alternative can be to apply XLA to the model forward pass only (as opposed to the whole `generate` method):
```python
import time
import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM

# 1. Load model and tokenizer
model_name = "distilgpt2"
# remember: decoder-only models need left-padding
tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
tokenizer.pad_token = tokenizer.eos_token
model = TFAutoModelForCausalLM.from_pretrained(model_name)

# 2. Prepare tokenization and generation arguments -- don't forget padding to avoid retracing!
tokenization_kwargs = {"pad_to_multiple_of": 32, "padding": True, "return_tensors": "tf"}
generation_kwargs = {"max_new_tokens": 32, "no_repeat_ngram_size": 5}

# 3. Replace the forward pass by the XLA equivalent
model.call = tf.function(model.call, jit_compile=True)

# 4. Generate! Remember -- the first call will be slow, but all subsequent calls will be fast if you've done things right.
input_prompts = [f"The best thing about {country} is" for country in ["Spain", "Japan", "Angola"]]
for input_prompt in input_prompts:
    tokenized_inputs = tokenizer([input_prompt], **tokenization_kwargs)
    start = time.time_ns()
    generated_text = model.generate(**tokenized_inputs, **generation_kwargs)
    end = time.time_ns()
    decoded_text = tokenizer.decode(generated_text[0], skip_special_tokens=True)
    print(f"Original prompt -- {input_prompt}")
    print(f"Generated -- {decoded_text}")
    print(f"Execution time -- {(end - start) / 1e6:.1f} ms\n")
```

Hi @gante, thank you so much for trying to make this operator XLA compatible! Due to that the generate method is currently the time bottleneck, apply XLA to the forward pass won't help much. We will see if we can figure out a solution for efficient implementation. 
I hope to get clarifications on this statement 'If it is not a hard requirement for your application, the easiest solution is not to use it for now :D (i.e. don't set the no_repeat_ngram_size argument)' --> Even if we don't set this no_repeat_ngram_size argument, all the inner calls still need to be compiled, so we still cannot compile the generate method, correct? I tried setting the argument to None it still cannot compile. 

Hey @mizhazha 👋 

If `no_repeat_ngram_size` is not set (or if it is set to `None`), then it will not be part of the code seen by the compiler, and XLA compilation should work.

What is the error you're seeing?

@gante Here is what I have run: 
````
from transformers import BartTokenizerFast, TFBartForConditionalGeneration
tokenizer = BartTokenizerFast.from_pretrained('facebook/bart-base')
xla_generate = tf.function(TFBartForConditionalGeneration.from_pretrained('facebook/bart-base').generate, jit_compile=True) 

tokenized_input = tokenizer("I have a white shiba", return_tensors="tf")
xla_generate(**tokenized_input)
````
Error message:
````
File "/home/jmi/miniconda3/envs/hf_tf29/lib/python3.7/site-packages/transformers/generation_tf_logits_process.py", line 427, in {}call{} * raise NotImplementedError("TFNoRepeatNGramLogitsProcessor is only implemented for eager execution.") NotImplementedError: TFNoRepeatNGramLogitsProcessor is only implemented for eager execution.
````

@mizhazha ah, that is because `generate` defaults are inherited from the model config file, and the model you are using sets this option ([here](https://huggingface.co/facebook/bart-base/blob/main/config.json#L44)). This inheritance is poorly documented at the moment, and we are working to improve it (https://github.com/huggingface/transformers/issues/18655).

Try setting `no_repeat_ngram_size=0` :) (`no_repeat_ngram_size=None` or not passing the argument causes the default from the model config to be used)

@gante Ah that works!! Thanks for the pointers and the help to unblock!

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.

Hi @gante , I am still facing below error XLA compatible error for [hugging-face summarization code](https://github.com/huggingface/notebooks/blob/main/examples/summarization-tf.ipynb).
```
_logits_process.py", line 428, in __call__  *
        raise NotImplementedError("TFNoRepeatNGramLogitsProcessor is only implemented for eager execution.")

    NotImplementedError: TFNoRepeatNGramLogitsProcessor is only implemented for eager execution.
```
I have Apple M1 Pro and I have installed latest transformer's git repo(`pip install --upgrade git+https://github.com/huggingface/transformers.git`). I read your above comments but don't know where to add this for summarization code.

Hey @kirtinikam 👋 This feature is not available on TF XLA and it's not on our plans to implement it. The solution is to turn it off, i.e. to call `.generate()` (or the pipeline) with `no_repeat_ngram_size=0`

Hi @gante,  I followed your suggestion and passed `no_repeat_ngram_size=0` in the `generate()` method. But still seeing XLA compilation errors on the Apple M1 Pro system.

```
  File "/Users/summarizer/venv/lib/python3.9/site-packages/tensorflow/compiler/tf2xla/ops/gen_xla_ops.py", line 1040, in xla_dynamic_slice
    return xla_dynamic_slice_eager_fallback(
  File "/Users/summarizer/venv/lib/python3.9/site-packages/tensorflow/compiler/tf2xla/ops/gen_xla_ops.py", line 1092, in xla_dynamic_slice_eager_fallback
    _result = _execute.execute(b"XlaDynamicSlice", 1, inputs=_inputs_flat,
tensorflow.python.framework.errors_impl.UnimplementedError: Exception encountered when calling layer "SelfAttention" (type TFT5Attention).

Could not find compiler for platform METAL: NOT_FOUND: could not find registered compiler for platform METAL -- check target linkage [Op:XlaDynamicSlice]

```
I tired  to set no_repeat_ngram_size in both ways, (1) pass no_repeat_ngram_size=0 parameter to `generate()` method or (2) pass in kwargs dictionary `generation_kwargs = {"no_repeat_ngram_size": 0}` but got same error. 
Can you please take a look at my code below. 

   ```
 def evaluate(self):
        print(f"perform evaluation")
        self.load_model('./models/my_model')
        # generation_kwargs = {"no_repeat_ngram_size": 0}

        all_preds = []
        all_labels = []
        for batch, labels in tqdm(self.generation_dataset):
            predictions = self.model.generate(input_ids=batch["input_ids"],
                                              attention_mask=batch["attention_mask"],
                                              max_new_tokens=32, no_repeat_ngram_size=0)
            decoded_preds = self.tokenizer.batch_decode(predictions, skip_special_tokens=True)
            labels = labels.numpy()
            labels = np.where(labels != -100, labels, self.tokenizer.pad_token_id)
            decoded_labels = self.tokenizer.batch_decode(labels, skip_special_tokens=True)
            decoded_preds = ["\n".join(sent_tokenize(pred.strip())) for pred in decoded_preds]
            decoded_labels = ["\n".join(sent_tokenize(label.strip())) for label in decoded_labels]
            all_preds.extend(decoded_preds)
            all_labels.extend(decoded_labels)

        result = self.metric.compute(
            predictions=all_preds, references=all_labels, use_stemmer=True
        )
        result = {key: value.mid.fmeasure * 100 for key, value in result.items()}
        print({k: round(v, 4) for k, v in result.items()})
```

Any sugestions on this.

Hey @kirtinikam -- that seems like a MacOS issue. Check [this link](https://developer.apple.com/forums/thread/709569) :)