[BugFix] Fix non-deterministic key order in stack (#1230)

(cherry picked from commit c35d7aa)

Author: Vincent Moens

tensordict/_torch_func.py

@@ -626,7 +626,6 @@ def stack_fn(key, values, is_not_init, is_tensor):
                 key: stack_fn(key, values, is_not_init, is_tensor)
                 for key, (values, is_not_init, is_tensor) in out.items()
             }
-
             result = clz._new_unsafe(
                 out,
                 batch_size=LazyStackedTensorDict._compute_batch_size(

tensordict/utils.py

@@ -1757,7 +1757,7 @@ def _check_keys(
     strict: bool = False,
     include_nested: bool = False,
     leaves_only: bool = False,
-) -> set[str]:
+) -> set[str] | list[str]:
     from tensordict.base import _is_leaf_nontensor
 
     if not len(list_of_tensordicts):
@@ -1769,27 +1769,29 @@ def _check_keys(
     )
     # TODO: compile doesn't like set() over an arbitrary object
     if is_compiling():
-        keys = {k for k in keys}  # noqa: C416
+        keys_set = {k for k in keys}  # noqa: C416
     else:
-        keys: set[str] = set(keys)
+        keys_set: set[str] = set(keys)
     for td in list_of_tensordicts[1:]:
         k = td.keys(
             include_nested=include_nested,
             leaves_only=leaves_only,
             is_leaf=_is_leaf_nontensor,
         )
         if not strict:
-            keys = keys.intersection(k)
+            keys_set = keys_set.intersection(k)
         else:
             if is_compiling():
                 k = {v for v in k}  # noqa: C416
             else:
                 k = set(k)
-            if k != keys:
+            if k != keys_set:
                 raise KeyError(
                     f"got keys {keys} and {set(td.keys())} which are incompatible"
                 )
-    return keys
+    if strict:
+        return list(keys)
+    return keys_set
 
 
 def _set_max_batch_size(source: T, batch_dims=None):

test/test_tensorclass.py

@@ -1884,6 +1884,20 @@ class MyDataNested:
         ):
             torch.stack([data1, data3], dim=0)
 
+    def test_stack_keyorder(self):
+
+        class MyTensorClass(TensorClass):
+            foo: Tensor
+            bar: Tensor
+
+        tc1 = MyTensorClass(foo=torch.zeros((1,)), bar=torch.ones((1,)))
+
+        for _ in range(10000):
+            assert list(torch.stack([tc1, tc1], dim=0)._tensordict.keys()) == [
+                "foo",
+                "bar",
+            ]
+
     def test_statedict_errors(self):
         @tensorclass
         class MyClass: