fix: Restore explicit .keys() calls for TensorDict compatibility
Author: pankajbaid567
Pull request: yes

# Fix: Restore explicit .keys() calls for TensorDict compatibility

Fixes: PreTrainedTokenizerBase.pad() broken on tensordict.TensorDict #42370
Open

## What does this PR do?

This PR restores TensorDict compatibility that was broken by PR #37283, which removed explicit `.keys()` calls from dictionary iteration patterns. The issue caused `RuntimeError: generator raised StopIteration` when using TensorDict objects with data collators and tokenization utilities.

### Problem

TensorDict objects implement `__iter__()` to iterate over batch dimensions rather than dictionary keys (as standard dicts do). When PR #37283 replaced explicit `.keys()` calls with implicit iteration (e.g., `for key in dict[0]`), it broke compatibility with TensorDict and similar Mapping implementations.

**Error encountered:**
```python
RuntimeError: generator raised StopIteration

```

**This affected:**

[DataCollatorWithPadding](vscode-file://vscode-app/private/var/folders/3r/brgdvkx91hv16l5vj24sd7k00000gn/T/AppTranslocation/11644A3F-AD49-4E85-B517-6AE093C8E246/d/Visual%20Studio%20Code%20-%20Insiders.app/Contents/Resources/app/out/vs/code/electron-browser/workbench/workbench.html)
[DataCollatorForLanguageModeling](vscode-file://vscode-app/private/var/folders/3r/brgdvkx91hv16l5vj24sd7k00000gn/T/AppTranslocation/11644A3F-AD49-4E85-B517-6AE093C8E246/d/Visual%20Studio%20Code%20-%20Insiders.app/Contents/Resources/app/out/vs/code/electron-browser/workbench/workbench.html)
[tokenizer.pad()](vscode-file://vscode-app/private/var/folders/3r/brgdvkx91hv16l5vj24sd7k00000gn/T/AppTranslocation/11644A3F-AD49-4E85-B517-6AE093C8E246/d/Visual%20Studio%20Code%20-%20Insiders.app/Contents/Resources/app/out/vs/code/electron-browser/workbench/workbench.html) method
Any code path that converts a list of dict-like objects into a dict of lists
Solution
Restored explicit .keys() calls in 5 critical locations where dict-list conversion happens during padding operations. This ensures compatibility with any object implementing the [Mapping](vscode-file://vscode-app/private/var/folders/3r/brgdvkx91hv16l5vj24sd7k00000gn/T/AppTranslocation/11644A3F-AD49-4E85-B517-6AE093C8E246/d/Visual%20Studio%20Code%20-%20Insiders.app/Contents/Resources/app/out/vs/code/electron-browser/workbench/workbench.html) protocol, regardless of how it implements __iter__().

**Pattern applied:**
```
# Before (broken for TensorDict):
encoded_inputs = {key: [example[key] for example in encoded_inputs] for key in encoded_inputs[0]}

# After (works for all Mapping objects):
encoded_inputs = {key: [example[key] for example in encoded_inputs] for key in encoded_inputs[0].keys()}
```
Also added a length check to prevent IndexError on empty batches:
`if isinstance(encoded_inputs, (list, tuple)) and len(encoded_inputs) > 0 and isinstance(encoded_inputs[0], Mapping):`
**Before submitting**
 

- [x] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).

 

- [x] Did you read the [contributor guideline](vscode-file://vscode-app/private/var/folders/3r/brgdvkx91hv16l5vj24sd7k00000gn/T/AppTranslocation/11644A3F-AD49-4E85-B517-6AE093C8E246/d/Visual%20Studio%20Code%20-%20Insiders.app/Contents/Resources/app/out/vs/code/electron-browser/workbench/workbench.html), Pull Request section?

 

- [x] Was this discussed/approved via a Github issue or the [forum](vscode-file://vscode-app/private/var/folders/3r/brgdvkx91hv16l5vj24sd7k00000gn/T/AppTranslocation/11644A3F-AD49-4E85-B517-6AE093C8E246/d/Visual%20Studio%20Code%20-%20Insiders.app/Contents/Resources/app/out/vs/code/electron-browser/workbench/workbench.html)? Please add a link to it if that's the case.

 

- [x] Did you make sure to update the documentation with your changes? (Added inline comments explaining the fix)

- [x]  Did you write any new necessary tests? (Yes - comprehensive test suite with 7 test cases)

**Who can review?**
@ArthurZucker @Rocketknight1 - This affects tokenizers and data collators @SunMarc - This affects the trainer data collators
**Additional Context:**

This is a minimal, surgical fix that only touches the specific lines where dict-list conversion happens. The fix is defensive and generalizes well to any Mapping implementation that might have non-standard __iter__() behavior.

Comments:
Hey, thank you for the PR! As far as I know, we aren't using `tensordict` anywhere in the codebase, so I don't know if we want tensordict-specific tests. Can we drop the tests and the `require_tensordict` / `is_tensordict_available` code, but keep the extra `.keys()`? The comment should just explain that it needs to stay for `tensordict` rather than referring to the issue number, since that forces users to look it up on Github to understand why it's there.

> Hey, thank you for the PR! As far as I know, we aren't using `tensordict` anywhere in the codebase, so I don't know if we want tensordict-specific tests. Can we drop the tests and the `require_tensordict` / `is_tensordict_available` code, but keep the extra `.keys()`?

In my humble opinion it's good to keep some minimal tests against `TensorDict` to avoid breaking a not-uncommon use case, despite no internal use or official support. Consider my opinion non-blocking though, maintainer’s judgment should take precedence.

> The comment should just explain that it needs to stay for `tensordict` rather than referring to the issue number, since that forces users to look it up on Github to understand why it's there.

Agreed, I take back my [previous comment](https://github.com/huggingface/transformers/pull/42373#discussion_r2558962140).

The failing test [test_training_gradient_checkpointing_use_reentrant_false](vscode-file://vscode-app/private/var/folders/3r/brgdvkx91hv16l5vj24sd7k00000gn/T/AppTranslocation/11644A3F-AD49-4E85-B517-6AE093C8E246/d/Visual%20Studio%20Code%20-%20Insiders.app/Contents/Resources/app/out/vs/code/electron-browser/workbench/workbench.html) in Longformer is unrelated to this PR. Here's why:

Scope of changes: This PR only modifies tokenization padding logic (adding explicit .keys() calls) in 5 tokenizer files. No changes to model architecture or gradient computation.

Luke's gradient checkpointing tests are already skipped: The Luke model (one of the modified tokenizers) has all three gradient checkpointing tests explicitly disabled with @unittest.skip because "This architecture seem to not compute gradients properly when using GC" (see [test_modeling_luke.py](vscode-file://vscode-app/private/var/folders/3r/brgdvkx91hv16l5vj24sd7k00000gn/T/AppTranslocation/11644A3F-AD49-4E85-B517-6AE093C8E246/d/Visual%20Studio%20Code%20-%20Insiders.app/Contents/Resources/app/out/vs/code/electron-browser/workbench/workbench.html)).

No connection between changes and failure:

Tokenization happens during data preprocessing
Gradient checkpointing happens during model forward/backward pass
These are completely separate phases with no interaction

@ligz08 @Rocketknight1 any more changes required ?


Hmn, I'd still remove the tests and `is_tensordict_available`! It's not a dependency anywhere, so our CI is unlikely to ever have it installed. As a result, those tests are just going to get skipped anyway.

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: luke, mluke

@Rocketknight1 I have made the required changes sir 


The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_42373). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.