TF generate refactor - XLA sample
# What does this PR do?

This PR brings XLA to `sample`, in `generate`. Four important details before reviewing:
1. The diff has the changes of https://github.com/huggingface/transformers/pull/16704, review that PR first plz :) It fixes a test from `beam_search`. I will rebase as soon as the other PR gets merged (the changes were bundled to confirm that it passes all generate tests).
2. The body is mostly copy/paste from `greedy_search`;
3. The sample step was changed from the previous implementation -- if we want to seed sampling with XLA, we need to use the `stateless` functions;
4. The XLA sample tests do not compare all generated tokens to their non-XLA sample counterparts, due to the numerical instabilities discussed on Slack. We do compare the first tokens, which are the same.

Finally, tests have been run for the usual models (`gpt2`, `t5`, `rag`, `speech2text`, `encoder_decoder`, `vision_encoder_decoder`, `bart`). 

____________________________

I've also run a quick sanity check on GPU. Using GPT2+sample, on an Nvidia T4:
- eager TF: ~1.7s
- XLA TF: ~54ms (~22s compile time) :point_right: 31x speedup
_The documentation is not available anymore as the PR was closed or merged._<|||||>@patrickvonplaten the `stateless` TF functions accept a `seed` argument that is a tuple of two integers 😅 Not very intuitive, I agree. They correspond to the `key` and `counter` used in the internal RNG algorithms ([source](https://www.tensorflow.org/api_docs/python/tf/random/Generator#from_key_counter)).

If you think it will be unintuitive for users, I can change it so that our `seed` argument corresponds to the `key` of the tuple (i.e. a single integer), and fix the `counter` to `0`. For practical purposes, it should be the same thing.<|||||>> @patrickvonplaten the `stateless` TF functions accept a `seed` argument that is a tuple of two integers sweat_smile Not very intuitive, I agree. They correspond to the `key` and `counter` used in the internal RNG algorithms ([source](https://www.tensorflow.org/api_docs/python/tf/random/Generator#from_key_counter)).
> 
> If you think it will be unintuitive for users, I can change it so that our `seed` argument corresponds to the `key` of the tuple (i.e. a single integer), and fix the `counter` to `0`. For practical purposes, it should be the same thing.

I see - ok maybe better to leave as is then to be aligned with TF<|||||>While running tests for T5 (as suggested by @Rocketknight1), I found out that our XLA code is not behaving properly for T5, for both `sample` and `greedy_search`. Because the problem is not exclusive to `sample`, I'm merging this PR and fixing the issue in a future one.

(example)
![image](https://user-images.githubusercontent.com/12240844/163792042-d23d01be-c2c9-40d4-bee2-30ea32c47a45.png)
