EncoderDecoder Cross Attention Generation Output Shape does not match Documentation
Author: l-salewski
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.5.1
- Platform: Linux-4.15.0-143-generic-x86_64-with-glibc2.27
- Python version: 3.9.4
- PyTorch version (GPU?): 1.7.1 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: Yes (v100)
- 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.
 -->
- encoderdecoder/text generation: @patrickvonplaten, @patil-suraj

## Information

Model I am using: EncoderDecoder with BERT

The problem arises when using:
* [x] the official example scripts: slightly modified/extended, see below

The tasks I am working on is:
* [x] my own task or dataset: just an example sentence from the docs

## To reproduce

Steps to reproduce the behavior:

1. Start with the example script from the [EncoderDecoder forward documenation](https://huggingface.co/transformers/model_doc/encoderdecoder.html#transformers.EncoderDecoderModel.forward)
2. Remove model training and model saving and loading steps (not relevant) and configure model to return attentions
3. Check shapes of cross attention outputs of generation and forward

```python
from transformers import EncoderDecoderModel, BertTokenizer
import torch

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = EncoderDecoderModel.from_encoder_decoder_pretrained('bert-base-uncased', 'bert-base-uncased') # initialize Bert2Bert from pre-trained checkpoints


# forward
input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
outputs = model(input_ids=input_ids, decoder_input_ids=input_ids, output_attentions=True)
forward_cross_attentions = outputs.cross_attentions

# As described in the docs the shapes are:
# "Tuple of torch.FloatTensor (one for each layer) of shape (batch_size, num_heads, sequence_length, sequence_length)"" (here (1,12,8,8))
print(f"Elements in forward cross attention: {len(forward_cross_attentions)}")
# Yields: Elements in forward cross attention: 12
print(f"Shapes in forward cross attention: {[fca.shape for fca in forward_cross_attentions]}")
# Yields: Shapes in forward cross attention: [torch.Size([1, 12, 8, 8]), torch.Size([1, 12, 8, 8]), torch.Size([1, 12, 8, 8]), torch.Size([1, 12, 8, 8]), torch.Size([1, 12, 8, 8]), torch.Size([1, 12, 8, 8]), torch.Size([1, 12, 8, 8]), torch.Size([1, 12, 8, 8]), torch.Size([1, 12, 8, 8]), torch.Size([1, 12, 8, 8]), torch.Size([1, 12, 8, 8]), torch.Size([1, 12, 8, 8])]


# generation
generated = model.generate(input_ids, decoder_start_token_id=model.config.decoder.pad_token_id, return_dict_in_generate=True, output_attentions=True)
generated_cross_attentions = generated.cross_attentions

# generated_cross_attentions contains 19 elements, maybe one for each generation step (generated.sequences has 20 elements)?
print(f"Elements in generation cross attention: {len(generated_cross_attentions)}")
# Yields: Elements in generation cross attention: 19

# All of the contained cross attentions have shape (1,12,1,8)
for cross_attention in generated_cross_attentions:
    print(f"Shapes in generation cross attention: {[gca.shape for gca in cross_attention]}")
# Yields: Shapes in generation cross attention: [torch.Size([1, 12, 1, 8]), torch.Size([1, 12, 1, 8]), torch.Size([1, 12, 1, 8]), torch.Size([1, 12, 1, 8]), torch.Size([1, 12, 1, 8]), torch.Size([1, 12, 1, 8]), torch.Size([1, 12, 1, 8]), torch.Size([1, 12, 1, 8]), torch.Size([1, 12, 1, 8]), torch.Size([1, 12, 1, 8]), torch.Size([1, 12, 1, 8]), torch.Size([1, 12, 1, 8])] (repeated 19 times)
```

Furthermore, if `num_beams>1` all `num_beams*batch_size` cross attentions are returned even if `num_return_sequences == 1`.

```python
# continued from above...
# generation
generated = model.generate(input_ids, decoder_start_token_id=model.config.decoder.pad_token_id, return_dict_in_generate=True, output_attentions=True, num_beams=10)
generated_cross_attentions = generated.cross_attentions

# generated_cross_attentions contains 19 elements, maybe one for each generation step (generated.sequences has 20 elements)?
print(f"Elements in generation cross attention: {len(generated_cross_attentions)}")

# All of the contained cross attentions have shape (10,12,1,8)
for cross_attention in generated_cross_attentions:
    print(f"Shapes in generation cross attention: {[gca.shape for gca in cross_attention]}")

print(f"Shape of the generated sequences: {generated.sequences.shape}")
Yields: torch.Size([1, 20])
```

## Expected behavior

A `Tuple of torch.FloatTensor (one for each layer) of shape (batch_size, num_heads, gen_sequence_length, sequence_length)`. Ideally the cross attentions batch size should match the batch size of the `generated.sequences`.

## Work arounds

Stacking the tuples and then concatting along the dimension which is 1 like this:
```python
torch.cat([torch.stack(ca) for ca in generated_cross_attentions], dim=-2)
```
yields such a tensor of the correct shape, is that the correct way to assemble it?

For the batch size issue I haven't found a work around yet. Is it possible to retain the beam indices of the selected beams from `generate`? `output_scores` is no help, because it has the same shape as `generated.sequences`.

Any help, ideas or pointers how to work around this are highly appreciated.

Comments:
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.