decoder_hidden_states output inconsistent when generating with SpeechEncoderDecoder models
Author: valentinp72
Pull request: no

### System Info

- `transformers` version: 4.23.1
- Platform: Linux-4.18.0-372.36.1.el8_6.x86_64-x86_64-with-glibc2.28
- Python version: 3.9.12
- Huggingface_hub version: 0.10.1
- PyTorch version (GPU?): 1.12.1+cu113 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes (a100)
- Using distributed or parallel set-up in script?: yes, but script below produces the same thing with one GPU

### Who can help?

@sanchit-gandhi @gante

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

When decoding a `SpeechEncoderDecoderModel`, we can output the decoder hidden states with the args `output_hidden_states` and `return_dict_in_generate` to `True`.
However, the output lengths are not consistent between the beam search decoding and greedy decoding, and between the output sequences and decoder hidden states.

The following example is using `wav2vec2-xls-r-300m-en-to-15`, and shows this inconsistency:

```python
from transformers import Wav2Vec2Processor, SpeechEncoderDecoderModel
import numpy as np

folder = "" # custom folder where the model are stored since we don't have internet on the compute nodes
device = 'cuda:0'

model = SpeechEncoderDecoderModel.from_pretrained(folder + "facebook/wav2vec2-xls-r-300m-en-to-15").to(device)
processor = Wav2Vec2Processor.from_pretrained(folder + "facebook/wav2vec2-xls-r-300m-en-to-15")

# generating dummy data with variable lengths
audios = [
    np.random.random((17_000 + i,))
    for i in range(10) # bsz of 10
]
input_values = processor(
    audios, return_tensors="pt", padding=True, sampling_rate=16_000,
).input_values.to(device)


# common parameters to both greedy and beam search decoding
common_params = {
    'max_new_tokens': 200,
    'output_hidden_states': True,
    'output_scores': True,
    'return_dict_in_generate': True
}

####
print("Greedy decoding:")
generated_greedy = model.generate(
    input_values,
    num_beams=1,
    **common_params
)
print("  sequences shape: ", generated_greedy['sequences'].shape)
print("  decoder_hidden_states len: ", len(generated_greedy['decoder_hidden_states']))

####
print("Beam search decoding:")
generated_beam_search = model.generate(
    input_values,
    num_beams=2,
    **common_params
)
print("  sequences shape: ", generated_beam_search['sequences'].shape)
print("  decoder_hidden_states len: ", len(generated_beam_search['decoder_hidden_states']))
```

The output of that script is:
```
Greedy decoding:
  sequences shape:  torch.Size([10, 3])
  decoder_hidden_states len:  2
Beam search decoding:
  sequences shape:  torch.Size([10, 3])
  decoder_hidden_states len:  39
```

Following the documentation for [GreedySearchEncoderDecoderOutput](https://huggingface.co/docs/transformers/main/en/internal/generation_utils#transformers.generation.GreedySearchEncoderDecoderOutput.decoder_hidden_states) and [BeamSearchEncoderDecoderOutput](https://huggingface.co/docs/transformers/main/en/internal/generation_utils#transformers.generation.BeamSearchEncoderDecoderOutput.decoder_hidden_states):
- **greedy**: decoder_hidden_states (tuple(tuple(torch.FloatTensor)), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of torch.FloatTensor of shape (batch_size, generated_length, hidden_size).
- **beam search**: decoder_hidden_states (tuple(tuple(torch.FloatTensor)), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of torch.FloatTensor of shape (batch_size*num_beams*num_return_sequences, generated_length, hidden_size).

I think the length of the tuple (one element for each generated token) should be the same as the `sequence_length`.

*PS: it seems that when the audio is long enough so that the output is capped at `max_new_tokens`, the sequences length is `max_new_tokens + 1`, and the hidden states `max_new_tokens`.*

### Expected behavior

The sequences output should be the same length as the decoder_hidden_states length, like:

```
Greedy decoding:
  sequences shape:  torch.Size([10, 3])
  decoder_hidden_states len:  3
Beam search decoding:
  sequences shape:  torch.Size([10, 3])
  decoder_hidden_states len:  3
```

or

```
Greedy decoding:
  sequences shape:  torch.Size([10, 2])
  decoder_hidden_states len:  2
Beam search decoding:
  sequences shape:  torch.Size([10, 39])
  decoder_hidden_states len:  39
```

Comments:
Hey @valentinp72 👋 

I don't see a cause for a bug, but I am aware that our docstrings are in need of improvements :) 

Allow me to elaborate:
1. Greedy Search: your output (`sequences`) will be `[batch_size, generated_tokens + 1]`, and the `decoder_hidden_states` will have length `generated_tokens`. `sequences` has the `+1` because the output sequence contains the BOS token, which is set before the first forward pass of the model, so there are no `decoder_hidden_states` for that token.
2. Beam Search: Here it's trickier. In essence, beam search looks for candidate outputs until it hits a stopping condition. The candidate outputs can have fewer tokens than the total number of generation steps -- for instance, in an encoder-decoder text model, if your input is `How much is 2 + 2?` and the model generates as candidates `<BOS>4<EOS>` (3 tokens) and `<BOS>The answer is potato<EOS>` (for argument's sake, 6 tokens) before deciding to stop, you should see `sequences` with shape `[1, 3]` and `decoder_hidden_states` with length `5`, because `5` tokens were generated internally before settling on the 1st candidate.

Does it make more sense now? 🤗 

Oh I see!
Indeed, the docs could be improved :)

So, if I'd want to have the *real* hidden states associated with the first candidate (in beam search), eg. `<BOS>4<EOS>`, I would need to extract only the first 3 hidden states from `decoder_hidden_states`?

For example, that should be correct, assuming the `decoder_hidden_states` are sorted the same way the candidates are ordered? (ie. `[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]` for a batch size of `10` and a beam width of `2`.)
```python
bsz = generated_beam_search['sequences'].shape[0]
seqlen = generated_beam_search['sequences'].shape[1]
featdim = 1024
decoder_hidden_states = torch.zeros((bsz, seqlen, featdim)).to(device)

for i in range(seqlen):
    all_beams = generated_beam_search['decoder_hidden_states'][i][-1][:,0,:]
    # filtering the hidden states to only the first candidate for each sample
    only_first = all_beams.index_select(
        dim=0, index=torch.tensor([x for x in range(0, bsz * beam_width, beam_width)]).to(device)
    )
    decoder_hidden_states[:,i] = only_first
```

If I'm correct, I think an option to return the hidden states in a tensor format (instead of tuples of tensors) according the output candidates could be nice, for both greedy and beam search decoding. 

@valentinp72 

> So, if I'd want to have the real hidden states associated with the first candidate (in beam search), eg. `<BOS>4<EOS>`, I would need to extract only the first 3 hidden states from `decoder_hidden_states`?

Almost correct! The first 2 hidden states (one for `4`, another for `<EOS>`. `<BOS>` has no corresponding hidden states)

As for the exact methodology to extract the right values from `decoder_hidden_states` with beam search, the plot thickens 😅 It is the same problem as extracting the token-level scores from the `scores` output in beam search -- see [this function](https://github.com/huggingface/transformers/blob/42b60f8b02941b0c40c42e150a101eb372c3856e/src/transformers/generation/utils.py#L927) and its examples. If you replace `scores` by `decoder_hidden_states`, it should be very close to what you want.

In a nutshell, the index of the n-th output sequence in beam search changes over the course of its execution. The output sequence with index 0 may correspond to the sequence with index 1 at the 1st generation step, the sequence with index 5 at the 2nd generation step, and so on. `beam_indices` contains the index of each output sequence at each beam search step, from which you can de-scramble `decoder_hidden_states`.

Thank you.
I've adapted the `compute_transition_scores` to what I want, but I still have one doubt about the contents of `beam_indices`.

Here is an example generated by the beam search:
```
tensor([[ 0,  0,  0,  3,  3,  1,  0,  2,  1,  0,  3,  0,  0, -1],
        [ 5,  5,  5,  8,  6,  5,  5,  7,  6,  6,  6,  5,  5, -1],
        [10, 10, 10, 13, 12, 10, 10, 10, 10, 10, 12, 10, 10, -1],
        [15, 15, 15, 18, 17, 15, 15, 16, 16, 15, 17, 15, 15, -1],
        [20, 20, 20, 23, 21, 20, 22, 21, 21, 20, 20, 20,  0, -1],
        [25, 25, 25, 28, 27, 26, 26, 27, 27, 26, 29, 26, 25, -1],
        [30, 30, 30, 33, 32, 31, 30, 32, 32, 32, 32, 30, 30, -1],
        [35, 35, 35, 38, 38, 36, 35, 37, 37, 36, 38, 35, 35, -1],
        [40, 40, 40, 43, 42, 40, 40, 40, 40, 40, 42, 40, 40, -1],
        [45, 45, 45, 48, 49, 46, 46, 47, 46, 45, 47, 45, 45, -1]],
       device='cuda:0')
```

If we take the first row (first sequence), does it means that the correct hidden states for that sequence can be found at indexes `0, 0, 0, 3, 3, 1, 0, 2, 1, 0, 3, 0, 0`? If so, why, for the 5th sequence, we have the index `0` for the last hidden states? Shouldn't each beam be independent?

@valentinp72 

> If so, why, for the 5th sequence, we have the index 0 for the last hidden states? Shouldn't each beam be independent?

They should 👀 can you share the snippet that leads to those `beam_indices`? That may be a bug. 

Now I'm no longer able to reproduce this error.
I think it was due to a bug on my own, while implementing my function I might have executed it twice, leading to (some?) -1 being replaced by 0.

I'm closing this issue as it seems my function that extracts the hidden representations works. I'm sharing it below if others needs it:

```python
def extract_decoder_hidden_states(
    generate_output_dict,
    hidden_layer_idx=-1,
):
    """
    Extracts the decoder hidden states representation from
    GreedySearchEncoderDecoderOutput and BeamSearchEncoderDecoderOutput,
    associated with the `sequences` output.
    - generate_output_dict: output dict from the model.generate() method
      you should add the following arguments to generate:
        - output_hidden_states=True
        - output_scores=True
        - return_dict_in_generate=True
    - hidden_layer_idx: index of the layer to extract the representation from (-1 == last one)
    """
    greedy = isinstance(generate_output_dict, GreedySearchEncoderDecoderOutput)
    beamy  = isinstance(generate_output_dict, BeamSearchEncoderDecoderOutput)

    if greedy:
        # in greedy decoding, the beam_indices is not present, so we create one
        # where the first beam is always selected
        scores = generate_output_dict['scores']
        device = generate_output_dict['sequences'].device
        beam_indices = torch.arange(scores[0].shape[0]).view(-1, 1)
        beam_indices = beam_indices.expand(-1, len(scores)).to(device)
    elif beamy:
        if 'beam_indices' not in generate_output_dict:
            raise RuntimeError(
                "You should export the scores with output_scores=True when " \
                "calling extract_decoder_hidden_states with " \
                "BeamSearchEncoderDecoderOutput"
            )
        beam_indices = generate_output_dict['beam_indices'].clone()
    else:
        raise NotImplementedError(
            "extract_decoder_hidden_states only works with " \
            "GreedySearchEncoderDecoderOutput and BeamSearchEncoderDecoderOutput " \
            "output types."
        )

    # handling of the target length and preparing the masking for tokens
    # outside of that length
    beam_indices_mask = beam_indices < 0
    max_beam_length = (1 - beam_indices_mask.long()).sum(-1).max()
    beam_indices = beam_indices[:, :max_beam_length]
    beam_indices_mask = beam_indices_mask[:, :max_beam_length]
    beam_indices[beam_indices_mask] = 0

    seqlen = generate_output_dict['sequences'].shape[1] - 1

    # creating the output hidden_states representation in format:
    # [bsz * beam_width ; seqlen ; featdim]
    decoder_hidden_states = torch.stack([
        generate_output_dict['decoder_hidden_states'][i][hidden_layer_idx][:,0,:].index_select(
            dim=0,
            index=beam_indices[:,i] # reordering using the beam_indices
        )
        for i in range(seqlen)
    ]).transpose(0, 1)

    # setting to 0 the hidden_states were it doesn't make sense to have an output
    decoder_hidden_states[beam_indices_mask] = 0

    return decoder_hidden_states
```