Strange shape of Scores vector
Author: andreabac3
Pull request: no

### System Info

transformers==4.20.1
torch==1.11.0+cu113
Python 3.9.13


### Who can help?

@patrickvonplaten @Narsil @ola13 @gante 

### Information

- [X] The official example scripts

### Reproduction

Hello everyone,

I am using BART and I have enabled `return_scores = True` with `beam_size = 4`.
The shape of the Scores vector is `(seq_len, batch_size * beam_size * num_returned sequence, vocab_size)`.

I would like to know how I can subdivide the vector obtaining the shape
`(seq_len, batch_size, beam_size, num_returned_sequence, vocab_size)`.

Because at the moment it is not possible to map the input sentences in a batch with the output.

Kind regards,
Andrea

ps: I also opened a [blog post](https://discuss.huggingface.co/t/strange-shape-of-scores-vector/24780), I did not know which was the most appropriate place, sorry for the duplicate.

```python3
outputs = model.generate(
            input_ids=source_ids,
            max_length = 500,
            return_dict_in_generate=True, 
            output_scores=True, 
            num_beams=4
        )

print(len(outputs.scores)) # seq_len
print(outputs.scores[0].shape) # (batch_size * beam_size * num_returned sequence, vocab_size)
```


### Expected behavior

a score vector with a shape of (seq_len, batch_size, beam_size, num_returned_sequence, vocab_size)

Comments:
Hi @andreabac3 👋 

I apologize in advance. The documentation for that part of the codebase is poor at the moment, so it's completely understandable that you feel confused.

The first documentation issue is the shape of `outputs.scores[0].shape`, which is actually `(batch_size * beam_size, vocab_size)`. It contains the scores (logits) of each token for each beam at each step. However, on their own, these scores are not very helpful. 

The most common use case is to use this tensor to obtain the scores of the selected tokens for each output -- we have an undocumented function for that (second documentation issue, tracked in https://github.com/huggingface/transformers/issues/18616), [`compute_transition_beam_scores`](https://github.com/huggingface/transformers/blob/e0b825a8d03f50ed9dbf9fbbbb3b4fcf0b4e4b22/src/transformers/generation_utils.py#L876). You can call it as 
```python
model.compute_transition_beam_scores(outputs.sequences, outputs.scores, outputs.beam_indices)
```
and it returns a tensor of shape `(num_return_sequences, seq_len)`. For instance, the value as `[2, 6]` corresponds to score of the 7th token for the 3rd returned sequence. If you sum these scores across the `seq_len` axis and divide by the length of the sequence (and apply the appropriate `length_penalty` scaling, if needed), you will obtain back `output.sequences_scores`.

I hope this helps! We will be revisiting the documentation soon to make all this clear. Let me know if you have further questions 🤗 

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/main/CONTRIBUTING.md) are likely to be ignored.

Hi @gante 👋, 
sorry for the late answer.

Don't worry about the missing documentation; the function worked perfectly.

Thanks again for your support,
Andrea