TF: XLA beam search + most generation-compatible models are now also XLA-generate-compatible
# What does this PR do?

The much-awaited PR -- beam search is now XLA compatible. GPT2 is the only model with XLA beam search tests, more models will follow in subsequent PRs 🎊  Preliminary tests on my machine shows that XLA beam search on GPU is ~26x faster (greedy search and sample are ~30x faster).

Slow tests have been run for the usual generate models (gpt2, t5, rag, speech_to_text, encoder_decoder, vision_encoder_decoder, bart).

EDIT: I've also generalized a few functions, and now ALL models that are compatible with generate are also compatible with XLA generate (with a few exceptions, when the models have no past cache support)

__________________________________

A hard-earned lesson which is kinda obvious in hindsight: `if` branches can make the XLA compiler confused about variable shapes, tagging their shape as `<unknown>`, which in turn causes all sorts of exceptions. Out of curiosity, I tried replacing the `if` by `tf.cond`, but the `<unknown>` shape persisted (because the tensor could indeed have a different shape at tracing time, depending on the branch taken)




_The documentation is not available anymore as the PR was closed or merged._<|||||>Very cool! Can we try it out for at least on Encoder-Decoder architecture as well (just to know that this code holds true here)?<|||||>@patrickvonplaten @Rocketknight1 now with encoder-decoder tests, and ready for review -- I was working on it on a separate branch, so I've merged it into this one. Now, this PR standardizes the XLA model kwargs preparation, and most models can use the XLA functionality. Some models were incompatible for different reasons, so there is a new flag to gate XLA generation (and the flag is set in the problematic architectures).

Finally, I'm also considering adding a general test like `test_xla_generate_fast`, but with `@slow`, beam search, and >100 tokens. It will probably break for a few models (like T5), but at least we would be able to automatically track which models are reliable with XLA beam search -- WDYT?  <|||||>Note: as per the comment above, if this PR gets merged as it is, I will open an issue to track issues regarding XLA generation (relevant models failing fast tests, as well as models failing the slow tests)