TF: BART compatible with XLA generation
Author: gante
Pull request: yes

# What does this PR do?

Adds `position_ids` to `TFBart`, so that we can do generation with a padded past -- a requirement for XLA generation.

This PR was built on top of #17426 (so it will contain its diff until it gets merged), and is a requirement for #17458. 

🚨 Important notes:
1. **Review suggestion**: check the Bart file, then its test file. The other changes are either cosmetic changes (e.g. correcting comments) or the result of `make fix-copies` (several files have copies from Bart).
2. There are several failing tests, but it's intentional -- some models' `prepare_inputs_for_generation` were copied from Bart, but the models do not have the `position_ids` input. If the PR gets a positive review, I will propagate the changes to the affected models.

Comments:
_The documentation is not available anymore as the PR was closed or merged._

@ydshieh tagging you for TF review, as Matt is off and you are also familiar with generate :)

> @ydshieh tagging you for TF review, as Matt is off and you are also familiar with generate :)

Actually not very familiar, but would love to get more involved 😃. Thanks for tagging me!

Hey @ydshieh 👋 answering your questions:

> From the change in prepare_inputs_for_generation, both in the PR for TF-GPT2 and this PR, my understanding of the main change is that: we need to use (decoder) attention mask in order to calculate the correct position_ids for both left/right padding. And this is done using tf.math.cumsum. Do I understand these PR correctly?

Correct 👍 

> Why we need decoder_position_ids when past_key_values is passed?

In the original PT code and eager execution TF, the position ids can be obtained by default (i.e. when not explicitly passed) from the past length, as the past length corresponds to the next position id if there is no left padding. In FLAX and XLA TF, the past is zero-padded, so the past length is not the default position id. As such, it is dangerous to leave the default path active -- this path should only be used in generate anyways, and the updated generate passes the position ids. (The GPT2 should also get the same guard, to be safe!)

> 

OK, I might got it. The `past` sent to the model is the padded (on the right) version! (which is required by XLA to have a fixed shape during loop, right?)

Thank you @gante ! 

I didn't think it in a thorough way, but in `prepare_inputs_for_generation`, when we return the actual inputs to a model,

https://github.com/huggingface/transformers/blob/9089b7b95a1b12f19b65872323f13f1f68a6eaa7/src/transformers/models/bart/modeling_tf_bart.py#L1428

it seems to me that we could cut `past` to the actual (non-padded) version. And when the model returns `past`, in `_update_model_kwargs_for_xla_generation`, we just always pad on the right.

(of course, we need to pass the current length info. to `prepare_inputs_for_generation` if we want to do so)

- this will keep `model_kwargs["past"]` compatible with XLA
- the actual `past` to model is the same as before 
  - especially, it won't get `max_length - 1` as length, so we no longer have overhead due to the increasing length 
- it might make the logic a bit easier in  `_update_model_kwargs_for_xla_generation`

@gante I don't want to make you too busy. I will let you judge if this is a good idea, and even if it is, if we should change it now, or we can do it later. I know we want to publish our work soon!


> it seems to me that we could cut past to the actual (non-padded) version. 

I would love to do that, and it would be a great idea to simplify the code, but sadly XLA does not allow dynamic-sized slices (i.e. cutting `past` based on the current length or based on its non-zero values). I've had the same idea too, but then I came across this limitation (documented [here](https://github.com/huggingface/transformers/pull/17378#issuecomment-1133641201))😢 Sadly, we have to keep working with the full padded array everywhere when XLA is on.

Think we can move towards finishing this PR here :-)

@patrickvonplaten it is ready to merge -- would you like to make a final review, or can I merge the PR? :)