❗ Bug for compute_transition_scores in generation 
Author: Hannibal046
Pull request: no

### System Info

- `transformers` version: 4.30.0.dev0
- Platform: Linux-5.15.0-1038-azure-x86_64-with-glibc2.17
- Python version: 3.8.16
- Huggingface_hub version: 0.14.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.1+cu117 (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?: no
- Using distributed or parallel set-up in script?: no

### Who can help?

@gante

### Information

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

### Tasks

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

### Reproduction

```python
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

src_lang = 'hye_Armn'
trg_lang = 'eng_Latn'
tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M",src_lang='hye_Armn')
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
batch = ['Հավելյալ 300-ը վագոնների թիվը դարձնում է 1,300, որոնց նպատակն է թեթևացնել գերծանրաբեռնվածությունը:']

inputs = tokenizer(batch,return_tensors="pt")
outputs = model.generate(
    **inputs, 
    forced_bos_token_id=tokenizer.lang_code_to_id[trg_lang],
    max_length=100,
    return_dict_in_generate=True,
    output_scores=True,
    num_return_sequences = 5,
    num_beams = 5,
)
transition_scores = model.compute_transition_scores(outputs.sequences, outputs.scores, normalize_logits=True)
```

### Expected behavior

transition_scores are computed

Comments:
@gante @ArthurZucker Hi, would you mind taking some time to check this?

Sure! Can you share the full traceback of the error that you stumbled upon? 

Thanks for help. The full traceback is here:
```
RuntimeError                              Traceback (most recent call last)
Cell In[1], line 19
      9 inputs = tokenizer(batch,return_tensors="pt")
     10 outputs = model.generate(
     11     **inputs, 
     12     forced_bos_token_id=tokenizer.lang_code_to_id[trg_lang],
   (...)
     17     num_beams = 5,
     18 )
---> 19 transition_scores = model.compute_transition_scores(outputs.sequences, outputs.scores, normalize_logits=True)

File [/anaconda/envs/llmt/lib/python3.8/site-packages/transformers/generation/utils.py:1086](https://vscode-remote+ssh-002dremote-002b7b22686f73744e616d65223a2241313030227d.vscode-resource.vscode-cdn.net/anaconda/envs/llmt/lib/python3.8/site-packages/transformers/generation/utils.py:1086), in GenerationMixin.compute_transition_scores(self, sequences, scores, beam_indices, normalize_logits)
   1084 # 7. Define which indices contributed to scores
   1085 cut_idx = sequences.shape[-1] - max_beam_length
-> 1086 indices = sequences[:, cut_idx:] + beam_sequence_indices
   1088 # 8. Compute scores
   1089 transition_scores = scores.gather(0, indices)

RuntimeError: The size of tensor a (2) must match the size of tensor b (23) at non-singleton dimension 1
```

BTW, this RuntimeError doesn't always happen. For example, for input like this, this snippet works fine.
```
batch = ['«Մենք հիմա ունենք 4 ամսական մկներ, որոնք, նախկինում շաքարային դիաբետ ունենալով, այժմ չունեն այն,- ավելացրեց նա։»']
```

Hey @Hannibal046 👋  

As stated in the [docs](https://huggingface.co/docs/transformers/v4.30.0/en/main_classes/text_generation#transformers.GenerationMixin.compute_transition_scores), when `num_beams>1`, you need to pass the `beam_indices` argument to `compute_transition_scores()`. `beam_indices` is part of the output in generate with beam search.

Here's a working snippet: 
```py
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

src_lang = 'hye_Armn'
trg_lang = 'eng_Latn'
tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M", src_lang='hye_Armn')
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
batch = ['Հավելյալ 300-ը վագոնների թիվը դարձնում է 1,300, որոնց նպատակն է թեթևացնել գերծանրաբեռնվածությունը:']

inputs = tokenizer(batch, return_tensors="pt")
outputs = model.generate(
    **inputs,
    forced_bos_token_id=tokenizer.lang_code_to_id[trg_lang],
    max_length=100,
    return_dict_in_generate=True,
    output_scores=True,
    num_return_sequences = 5,
    num_beams = 5,
)
transition_scores = model.compute_transition_scores(outputs.sequences, outputs.scores, outputs.beam_indices, normalize_logits=True)
print(transition_scores)
```

@gante Thanks so much! Appreciate your work!

@gante we can maybe raise an error if indices are not properly passed? 

@ArthurZucker Sadly that is not possible without an API change :( We do get an exception when `num_beams` and `num_return_sequences` are used together in `generate`, but not when `num_beams` is used alone -- the output format is the same as a batched input, no way to detect whether it comes from beam search or not.

E.g. this snippet runs (and it should throw an error)
```py
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")

inputs = tokenizer(["The quick brown"], return_tensors="pt")
gen_out = model.generate(**inputs, num_beams=5, do_sample=False, return_dict_in_generate=True, output_scores=True)
transition_scores = model.compute_transition_scores(gen_out.sequences, gen_out.scores, normalize_logits=True)
```

The solution would be e.g. to pass the entire outputs of generate into `compute_transition_scores`. But that's a bigger change that implies a deprecation cycle (that I'm not sure is worth going through 🤔 )