[BUG] Random 'serialization' when calling keys() vs flatten_keys().keys()

[ludwigwinkler]

Describe the bug

The order in which key sequences via keys() and flatten_keys().keys() are generated differ (in some cases).

I need to extract the tensors from the tensordict with tensordict.values(True, True) and pass them into torch.autograd.grad as a tuple.
I kept running into tensordict-reconstruction issues in which keys and the tensors did not match anymore to previous incarnations of the "same" tensordict and subsequent operations were wrong.

Here, state.x and state.R, and gradients.state.x and gradients.state.R are generated differently.
A simply alpha-numerical ordering would be appreciably.

To Reproduce

This is hard as I still haven't figured out why this occurs in my case.
For all my base test cases, it behaves normally, but once I use a large model to test things, I keep running into this problem.
Nevertheless, I think it would be a useful addition.

System info

import tensordict, numpy, sys, torch
print(tensordict.__version__, numpy.__version__, sys.version, sys.platform, torch.__version__)

0.4.0 1.26.4 3.11.9 (main, Apr 19 2024, 11:43:47) [Clang 14.0.6 ] darwin 2.3.1

Reason and Possible fixes

Can we squeeze a sorted option to tensordict.items()?
Most tensordicts would have a small number of keys (<10), so maybe it would be useful as the default?