Repository: huggingface/transformers

Issue:
TF generate() function is incompatible with output_attention and output_hidden_states 
## Environment info
   
- `transformers` version: 3.4.0
- Platform: Mac OS Catalina (10.15.6)
- Python version: 3.6.8
- PyTorch version (GPU?):  N/A
- Tensorflow version (GPU?): 2.3.1 (no)
- Using GPU in script?: No, but bug is persistent regardless of device.
- Using distributed or parallel set-up in script?: No

### Who can help
 @sshleifer  @TevenLeScao  @patrickvonplaten
 
## Information

The generate() function in modeling_tf_utils assumes that outputs from a model call have a static number of outputs. If either `output_attention` or `output_hidden_states` is set, the number of outputs is different, causing the function to fail. The fix should be pretty simple and only involve checking for the output size (or completely switching to dict/NamedTuple outputs from model modules since variable length returns are brittle :)).

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [X] my own modified scripts: (give details below)

The tasks I am working on is:
* [X] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Create a TF model with one of the above mentioned flags set.
2. Call `.generate()` on the model.

```python
import transformers 
model = transformers.TFT5ForConditionalGeneration.from_pretrained('t5-small', output_hidden_states=True, output_attentions=True)
tokenizer = transformers.T5Tokenizer.from_pretrained('t5-small')
input_ids = tokenizer.batch_encode_plus(['test 1', 'test 2', 'test 3'], return_tensors="tf", padding='longest')
output_ids = model.generate(input_ids['input_ids'], attention_mask=input_ids['attention_mask'])

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/transformers/generation_tf_utils.py", line 405, in generate
    use_cache=use_cache,
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/transformers/generation_tf_utils.py", line 445, in _generate_no_beam_search
    outputs = self(**model_inputs)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py", line 985, in __call__
    outputs = call_fn(inputs, *args, **kwargs)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/transformers/modeling_tf_t5.py", line 1352, in call
    training=training,
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py", line 985, in __call__
    outputs = call_fn(inputs, *args, **kwargs)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/transformers/modeling_tf_t5.py", line 759, in call
    training=training,
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py", line 985, in __call__
    outputs = call_fn(inputs, *args, **kwargs)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/transformers/modeling_tf_t5.py", line 450, in call
    assert len(past_key_value) == expected_num_past_key_values, error_message
AssertionError: There should be 4 past states. 2 (past / key) for self attention.2 (past / key) for cross attention Got 3 past key / value states
```

## Expected behavior

This snippet should not crash and have the same behavior as the one below. Though one may argue that the `generate()` function in this case should also return states/attentions which would complicate things. However, even ignoring the flags when generating is better than crashing. 

```python
import transformers 
model = transformers.TFT5ForConditionalGeneration.from_pretrained('t5-small', output_hidden_states=False, output_attentions=False)
tokenizer = transformers.T5Tokenizer.from_pretrained('t5-small')
input_ids = tokenizer.batch_encode_plus(['test 1', 'test 2', 'test 3'], return_tensors="tf", padding='longest')
output_ids = model.generate(input_ids['input_ids'], attention_mask=input_ids['attention_mask'])
print(output_ids)

tf.Tensor(
[[   0 2300  209    1    0]
 [   0 2300  794  204    1]
 [   0 2300  220    1    0]], shape=(3, 5), dtype=int32)
```


Discussion:

Resolving patch:
diff --git a/src/transformers/generation_tf_utils.py b/src/transformers/generation_tf_utils.py
--- a/src/transformers/generation_tf_utils.py
+++ b/src/transformers/generation_tf_utils.py
@@ -348,8 +348,7 @@ def generate(
                 shape=(-1,),
             )
             # expand encoder_outputs
-            encoder_outputs = (tf.gather(encoder_outputs[0], expanded_batch_idxs, axis=0), *encoder_outputs[1:])
-
+            encoder_outputs = (tf.gather(encoder_outputs[0], expanded_batch_idxs, axis=0),)
         else:
             encoder_outputs = None
             cur_len = shape_list(input_ids)[-1]


Test patch: