Documents of `past_key_values` in input and output for `PegasusModel` are not aligned
Author: ryangawei
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.6.1
- Platform: Linux-4.15.0-144-generic-x86_64-with-glibc2.10
- Python version: 3.8.5
- PyTorch version (GPU?): 1.6.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:



Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

Model hub:

- for issues with a model report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->
Pegasus @patrickvonplaten, @patil-suraj

## Information

According to the documentation of `PegasusModel`, the `past_key_values` for input and output have different shape,

```
past_key_values (Tuple[Tuple[torch.Tensor]] of length config.n_layers with each tuple having 2 tuples each of which has 2 tensors of shape (batch_size, num_heads, sequence_length - 1, embed_size_per_head)) –

Contains precomputed key and value hidden-states of the attention blocks. Can be used to speed up decoding.
```
```
past_key_values (tuple(tuple(torch.FloatTensor)), optional, returned when use_cache=True is passed or when config.use_cache=True) – 

Tuple of tuple(torch.FloatTensor) of length config.n_layers, with each tuple having 2 tensors of shape (batch_size, num_heads, sequence_length, embed_size_per_head)) and 2 additional tensors of shape (batch_size, num_heads, encoder_sequence_length, embed_size_per_head).
```

I'm trying to reproduce the behaviour of passing `past_key_values` as inputs, so I construct a `dummy_past_key_values` to feed into the model,
```
decoder_seq_length = decoder_input_ids.shape[1]
dummy_past_value_keys = torch.ones(size=[1, model.config.num_attention_heads, decoder_seq_length-1, int(model.config.d_model / model.config.num_attention_heads)], dtype=torch.float32)

pkv_tuple = ((dummy_past_value_keys, dummy_past_value_keys), (dummy_past_value_keys, dummy_past_value_keys))
pkv_tuple = tuple([pkv_tuple] * model.config.num_hidden_layers)

outputs = model(input_ids, decoder_input_ids=decoder_input_ids, past_key_values=pkv_tuple)
```
Then I got the following error,

```
AttributeError: 'tuple' object has no attribute 'shape'
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-22-be142139ae01> in <module>
----> 1 outputs = model(input_ids, decoder_input_ids=decoder_input_ids, past_key_values=pkv_tuple)
~/anaconda3/envs/wga/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),
~/anaconda3/envs/wga/lib/python3.8/site-packages/transformers/models/pegasus/modeling_pegasus.py in forward(self, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, head_mask, decoder_head_mask, cross_attn_head_mask, encoder_outputs, past_key_values, inputs_embeds, decoder_inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict)
   1267                 )
   1268 
-> 1269         outputs = self.model(
   1270             input_ids,
   1271             attention_mask=attention_mask,
~/anaconda3/envs/wga/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),
~/anaconda3/envs/wga/lib/python3.8/site-packages/transformers/models/pegasus/modeling_pegasus.py in forward(self, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, head_mask, decoder_head_mask, cross_attn_head_mask, encoder_outputs, past_key_values, inputs_embeds, decoder_inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict)
   1151 
   1152         # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn)
-> 1153         decoder_outputs = self.decoder(
   1154             input_ids=decoder_input_ids,
   1155             attention_mask=decoder_attention_mask,
~/anaconda3/envs/wga/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),
~/anaconda3/envs/wga/lib/python3.8/site-packages/transformers/models/pegasus/modeling_pegasus.py in forward(self, input_ids, attention_mask, encoder_hidden_states, encoder_attention_mask, head_mask, cross_attn_head_mask, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict)
    941 
    942         # past_key_values_length
--> 943         past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
    944 
    945         if inputs_embeds is None:
AttributeError: 'tuple' object has no attribute 'shape'
```

And if I make a `dummy_past_value_keys` using the shape described by the output document,

```
pkv_tuple = (dummy_past_value_keys,) * 4
pkv_tuple = tuple([pkv_tuple] * model.config.num_hidden_layers)
outputs = model(input_ids, decoder_input_ids=decoder_input_ids, past_key_values=pkv_tuple)
```
No error shows up. Maybe I misunderstood the documentation, but from my perspective `Tuple of length config.n_layers with each tuple having 2 tuples each of which has 2 tensors` seems more like a `Tuple[Tuple[Tuple[torch.Tensor]]]`, and the description itself is quite confusing. The code that throws the exception `past_key_values_length = past_key_values[0][0].shape[2] ` tries to access the tensor's shape, which can be only done with `Tuple[Tuple[torch.Tensor]]`, and it makes more sense that the input and output `past_key_values` have the same shape. I wonder if the description of input `past_key_values` is the old version and hasn't been updated? Thank you.


Comments:
Hi @AlfredWGA , you are right ! The description of input `past_key_values` is old and should be updated. 
The correct shape is as described by the output docstring. Thanks for reporting!