TF: generate without tf.TensorArray

[gante]

What does this PR do?

Some models, like XLNet, need more than just the previous token when past is used. This PR solves this problem with the help of some refactoring -- we no longer use TensorArray, instead we scatter updates into a fixed-size tensor. This refactor simplifies generate, especially beam_search, which may prove to be helpful in enabling XLA.

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

Why was this refactor needed?

As it can be read in this issue, TensorArray is meant to be used as a write-once array, anything else falls in the unexpected behavior domain -- in other words, our use was dangerous. The original solution to the XLNet problem was to read all existing tokens from the TensorArray, using the same logic as in this PR, but it failed with XLA -- and the behavior depended on what was written into the variable on its first write. Since we use fixed-size tensors, a normal tensor works just fine, and with simpler code (assuming the reader is familiar with how scatter works :D ).

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[gante]

cc @ydshieh -- this PR fixes the XLNet generate error we have been seeing :)

[patrickvonplaten]

Cool!

[Rocketknight1]

LGTM! (And agree about TensorArray being cursed). Did you see any performance changes from doing it this way?

[gante]

@Rocketknight1 no differences in terms of execution speed 👍

GPT2 sample on a 3090, average of 10 runs (excluding compilation time)

• Eager: 884 ms -> 888 ms
• XLA: 29.2 ms -> 29.3 ms
• JAX: 19.7 ms