

TITLE
TF: XLA model output differs when certain outputs are passed

COMMENTS
24

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
Depending on the passed inputs, the output of an XLA-compiled model may significantly differ from its non-XLA counterpart. This suggests we should add tests for XLA-output equivalence, just like we do with e.g. PT-TF, as it is not guaranteed.

At the moment, this blocks further developments in `generate()` (can't reliably reproduce non-XLA results with XLA). I will assess this problem for T5 (first model where I've noticed this), then check whether it is present for other key models, and finally add equivalence tests.

cc @patrickvonplaten @Rocketknight1 (feel free to pitch in with ideas and suggestions)

____________________________________________________
Example for reproducibility (updated: assert diff < x-> print diff):
```python
import tensorflow as tf
from transformers import TFT5ForConditionalGeneration, T5Tokenizer


tokenizer = T5Tokenizer.from_pretrained("t5-base")
model = TFT5ForConditionalGeneration.from_pretrained("t5-base")
model_xla = tf.function(model, jit_compile=True)
pad_token_id = model.config.pad_token_id

sentence_1 = "Translate English to German: I have a cat, two dogs, three horses, and four birds."
sentence_2 = "Translate English to German: I have a cat, two dogs, and three horses."

ids_single = tokenizer([sentence_1], return_tensors="tf", padding=True).input_ids
decoder_ids_single = tf.zeros((1, 1), dtype=tf.int32)
attention_single = tf.cast(tf.math.not_equal(ids_single, pad_token_id), dtype=tf.int32)  # as computed in generate

ids_pair = tokenizer([sentence_1, sentence_2], return_tensors="tf", padding=True).input_ids
decoder_ids_pair = tf.zeros((2, 1), dtype=tf.int32)
attention_pair = tf.cast(tf.math.not_equal(ids_pair, pad_token_id), dtype=tf.int32)  # as computed in generate

# case 1: with batch size = 1 and NO attention mask, XLA and non-XLA match
outputs = model(input_ids=ids_single, decoder_input_ids=decoder_ids_single)
outputs_xla = model_xla(input_ids=ids_single, decoder_input_ids=decoder_ids_single)
print(tf.math.reduce_max(tf.math.abs(outputs.logits - outputs_xla.logits)).numpy())

# case 2: with batch size > 1 and NO attention mask, XLA and non-XLA match
outputs = model(input_ids=ids_pair, decoder_input_ids=decoder_ids_pair)
outputs_xla = model_xla(input_ids=ids_pair, decoder_input_ids=decoder_ids_pair)
print(tf.math.reduce_max(tf.math.abs(outputs.logits - outputs_xla.logits)).numpy())

# case 3 FAILING: with batch size = 1 and attention mask, XLA and non-XLA match
outputs = model(input_ids=ids_single, decoder_input_ids=decoder_ids_single, attention_mask=attention_single)
outputs_xla = model_xla(input_ids=ids_single, decoder_input_ids=decoder_ids_single, attention_mask=attention_single)
print(tf.math.reduce_max(tf.math.abs(outputs.logits - outputs_xla.logits)).numpy())

# case 4 FAILING: with batch size < 1 and attention mask, XLA and non-XLA match
outputs = model(input_ids=ids_pair, decoder_input_ids=decoder_ids_pair, attention_mask=attention_pair)
outputs_xla = model_xla(input_ids=ids_pair, decoder_input_ids=decoder_ids_pair, attention_mask=attention_pair)
print(tf.math.reduce_max(tf.math.abs(outputs.logits - outputs_xla.logits)).numpy())
```
