[BugFix] Consolidate lazy stacks of non-tensors

ghstack-source-id: d3d822d
Pull Request resolved: #1224

(cherry picked from commit f67a15c)

Author: Vincent Moens

tensordict/base.py

@@ -96,6 +96,7 @@
     lazy_legacy,
     lock_blocked,
     prod,
+    set_capture_non_tensor_stack,
     set_lazy_legacy,
     strtobool,
     TensorDictFuture,
@@ -9078,7 +9079,8 @@ def newfn(item_and_out):
                     from tensordict._lazy import LazyStackedTensorDict
 
                     # We want to be able to return whichever data structure
-                    out = LazyStackedTensorDict.maybe_dense_stack(imaplist, dim)
+                    with set_capture_non_tensor_stack(False):
+                        out = LazyStackedTensorDict.maybe_dense_stack(imaplist, dim)
                 else:
                     out = torch.cat(imaplist, dim)
             return out

test/test_tensordict.py

@@ -39,6 +39,8 @@
 
 from tensordict import (
     get_defaults_to_none,
+    lazy_legacy,
+    lazy_stack,
     LazyStackedTensorDict,
     make_tensordict,
     PersistentTensorDict,
@@ -65,7 +67,6 @@
     convert_ellipsis_to_idx,
     is_non_tensor,
     is_tensorclass,
-    lazy_legacy,
     logger as tdlogger,
     set_lazy_legacy,
 )
@@ -11190,6 +11191,11 @@ def test_map_iter_interrupt_early(self, chunksize, num_chunks, shuffle):
 
 
 class TestNonTensorData:
+    @tensorclass
+    class SomeTensorClass:
+        a: str
+        b: torch.Tensor
+
     @pytest.fixture
     def non_tensor_data(self):
         return TensorDict(
@@ -11204,6 +11210,27 @@ def non_tensor_data(self):
             batch_size=[],
         )
 
+    @set_capture_non_tensor_stack(False)
+    def test_consolidate_nested(self):
+        import pickle
+
+        td = TensorDict(
+            a=TensorDict(b=self.SomeTensorClass(a="a string!", b=torch.randn(10))),
+            c=TensorDict(d=NonTensorData("another string!")),
+        )
+        td = lazy_stack([td.clone(), td.clone()])
+        td = lazy_stack([td.clone(), td.clone()], -1)
+
+        tdc = td.consolidate()
+
+        assert (tdc == td).all()
+
+        tdr = pickle.loads(pickle.dumps(td))
+        assert (tdr == td).all()
+
+        tdcr = pickle.loads(pickle.dumps(tdc))
+        assert (tdcr == td).all()
+
     def test_comparison(self, non_tensor_data):
         non_tensor_data = non_tensor_data.exclude(("nested", "str"))
         assert (non_tensor_data | non_tensor_data).get_non_tensor(("nested", "bool"))