Issue: Unexpected Shape of logits:  When Using generate() with num_return_sequences > 1
Author: athmanar
Pull request: no
Labels: name: bug

Hello,

I'm encountering an issue with the `generate()` function when using the **LLaMA 3.1 8B** model and setting `num_return_sequences > 1`.

### 🔍 What I'm trying to do:

I want to generate multiple sequences (e.g., `num_return_sequences=2`) and retrieve the **logits for each sequence individually**. Specifically, I expect:

- A list of size equal to `num_return_sequences`
- Each element in the list to be a tensor of shape `[NUM_TOKENS_GENERATED_IN_THIS_SEQUENCE, VOCAB_SIZE]`
- NUM_TOKENS_GENERATED_IN_THIS_SEQUENCE is variable [example 10 tokens in first sequence and 20 tokens in second sequence]

### 🧪 Code

Here’s how I’m calling `generate`:

```python
kwargs["return_dict_in_generate"] = True
kwargs["output_scores"] = True
kwargs["output_logits"] = True

out = self.model.generate(
    input_tokens,
    attention_mask=attention_mask,
    max_new_tokens=self.max_new_tokens,
    do_sample=do_sample,
    num_return_sequences=num_samples,
    pad_token_id=self.tokenizer.eos_token_id,
    stop_strings=self.eos,
    tokenizer=self.tokenizer,
    **kwargs,
)

logits = out.logits
```

❗ Problem

The returned logits is a tuple of length N, where each entry is a tensor of shape [num_return_sequences, vocab_size].

**This doesn’t seem correct or helpful in the case where each generated sequence can have a different number of tokens (due to sampling, early stopping, etc.). I suspect padding is applied internally, but it's unclear how to map each logit back to its corresponding sequence and token.**

Instead, it would make sense for the output to be a list of tensors, where each tensor corresponds to one generated sequence and has shape: [NUM_TOKENS_GENERATED_IN_THIS_SEQUENCE, VOCAB_SIZE]

✅ Hacky Workaround (Works, But Slow)
If I loop manually and generate one sequence at a time, the returned logits are correct:


```
ret = []
for jj in range(num_samples):
    out = self.model.generate(
        input_tokens,
        attention_mask=attention_mask,
        max_new_tokens=self.max_new_tokens,
        do_sample=do_sample,
        # num_return_sequences=num_samples,  <-- intentionally commented
        pad_token_id=self.tokenizer.eos_token_id,
        stop_strings=self.eos,
        tokenizer=self.tokenizer,
        **kwargs,
    )
    ret.append(out.logits)
```
This works (because the generation length varies naturally), but it is slow and inefficient, and defeats the purpose of num_return_sequences.

🙏 Request
Could you please clarify:

How can I correctly parse the logits from generate() when using num_return_sequences > 1?
Shouldn't the API return something more structured (like a list of per-sequence logits)?

Is this a known issue or is there a better way?
cc: @zucchini-nlp  @gante


### Who can help?

_No response_

### Information

- [x] The official example scripts
- [ ] 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

kwargs["return_dict_in_generate"] = True
kwargs["output_scores"] = True
kwargs["output_logits"] = True

out = self.model.generate(
    input_tokens,
    attention_mask=attention_mask,
    max_new_tokens=self.max_new_tokens,
    do_sample=do_sample,
    num_return_sequences=num_samples,
    pad_token_id=self.tokenizer.eos_token_id,
    stop_strings=self.eos,
    tokenizer=self.tokenizer,
    **kwargs,
)

### Expected behavior

not a padded logits of same shape.

Comments:
@athmanar hey!

The returned logits usually have shape `[1, vocab_size]` per generated token, which also includes pad tokens. In case if one seq is shorter than the other, it is padded on the right, and the pad token is fed to the model. Thus when `num_return_sequences=2`, you get ` [num_return_sequences, vocab_size]` per generated token and can discard `pad` tokens from analysis by matching with generated tokens

There is also a helper function that you can use to get probabilities per generated token, more info [in this forum thread](https://discuss.huggingface.co/t/announcement-generation-get-probabilities-for-generated-output/30075)

@zucchini-nlp hey thanks for pointing it out. For Anyone facing this issue you can do something like this:

```
             scores=generated_tokens['logits']
             L=promt_inputs['input_ids'].shape[-1]
            mask=(generated_tokens['sequences'][:,L:]!=tokenizer.pad_token_id)
            mask = mask.bool() #not sure if it is bool just to make sure.
            masked_scores = [batched_Scores[i, mask[i]] for i in range(batched_Scores.size(0))]
```

This will ensure than the logits of padded tokens are removed /masked and you will be left with a list of lists of of shape[ [Token1, vocab_size] , [Token2, vocab_size].... ]

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.