[BugFix] Ensure that maybe_dense_stack preserves the TC type

ghstack-source-id: 8972977
Pull Request resolved: #1252

Author: Vincent Moens

tensordict/_torch_func.py

@@ -464,6 +464,7 @@ def _stack(
         return type(list_of_tensordicts[0])._stack_non_tensor(
             list_of_tensordicts, dim=dim
         )
+    list_of_tensordicts_orig = list_of_tensordicts
     if is_tc:
         list_of_tensordicts = [tc._tensordict for tc in list_of_tensordicts]
         clz = type(list_of_tensordicts[0])
@@ -503,7 +504,7 @@ def _stack(
                     if maybe_dense_stack:
                         with set_lazy_legacy(True):
                             return _stack(
-                                list_of_tensordicts,
+                                list_of_tensordicts_orig,
                                 dim=dim,
                                 maybe_dense_stack=maybe_dense_stack,
                             )
@@ -537,7 +538,7 @@ def _stack(
                             lazy_stack_dim += 1
                         else:
                             dim = dim - 1
-                        return LazyStackedTensorDict(
+                        result = LazyStackedTensorDict(
                             *[
                                 _stack(
                                     list(subtds),
@@ -550,12 +551,16 @@ def _stack(
                             ],
                             stack_dim=lazy_stack_dim,
                         )
+                        if is_tc:
+                            return clz._from_tensordict(result)
+                        return result
+
                 lazy_stack_dim = list_of_tensordicts[0].stack_dim
                 if dim <= lazy_stack_dim:
                     lazy_stack_dim += 1
                 else:
                     dim = dim - 1
-                return LazyStackedTensorDict(
+                result = LazyStackedTensorDict(
                     *[
                         _stack(list_of_td, dim, maybe_dense_stack=maybe_dense_stack)
                         for list_of_td in _zip_strict(
@@ -564,6 +569,9 @@ def _stack(
                     ],
                     stack_dim=lazy_stack_dim,
                 )
+                if is_tc:
+                    return clz._from_tensordict(result)
+                return result
 
             out = {}
             for key in keys:
@@ -594,7 +602,7 @@ def _stack(
                                 if maybe_dense_stack:
                                     with set_lazy_legacy(True):
                                         return _stack(
-                                            list_of_tensordicts,
+                                            list_of_tensordicts_orig,
                                             dim=dim,
                                             maybe_dense_stack=maybe_dense_stack,
                                         )
@@ -641,6 +649,9 @@ def stack_fn(key, values, is_not_init, is_tensor):
                 *list_of_tensordicts,
                 stack_dim=dim,
             )
+            if is_tc:
+                return td_types[0]._from_tensordict(out)
+            return out
     else:
         keys = _check_keys(list_of_tensordicts)
         batch_size = list(batch_size)

test/test_tensorclass.py

@@ -1889,7 +1889,8 @@ class MyDataNested:
         assert squeeze_tc.y.X.shape == torch.Size([4, 5])
         assert squeeze_tc.z == squeeze_tc.y.z == z
 
-    @pytest.mark.parametrize("lazy", [True, False])
+    @set_capture_non_tensor_stack(False)
+    @pytest.mark.parametrize("lazy", [True, False, "maybe"])
     def test_stack(self, lazy):
         @tensorclass
         class MyDataNested:
@@ -1898,23 +1899,40 @@ class MyDataNested:
             y: "MyDataNested" = None
 
         X = torch.ones(3, 4, 5)
+        if lazy:
+            Xb = torch.randn(3, 4, 4)
+        else:
+            Xb = X.clone()
         z = "test_tensorclass"
         batch_size = [3, 4]
         data_nest = MyDataNested(X=X, z=z, batch_size=batch_size)
+        data_nest_b = MyDataNested(X=Xb, z=z, batch_size=batch_size)
         data1 = MyDataNested(X=X, y=data_nest, z=z, batch_size=batch_size)
-        data2 = MyDataNested(X=X, y=data_nest, z=z, batch_size=batch_size)
+        data2 = MyDataNested(X=Xb, y=data_nest_b, z=z, batch_size=batch_size)
 
-        if lazy:
+        if lazy is True:
             stacked_tc = LazyStackedTensorDict.lazy_stack([data1, data2], 0)
+        elif lazy == "maybe":
+            stacked_tc = LazyStackedTensorDict.maybe_dense_stack([data1, data2], 0)
         else:
             with set_capture_non_tensor_stack(True):
                 stacked_tc = torch.stack([data1, data2], 0)
         assert type(stacked_tc) is type(data1)
         assert isinstance(stacked_tc.y, type(data1.y))
-        assert stacked_tc.X.shape == torch.Size([2, 3, 4, 5])
-        assert stacked_tc.y.X.shape == torch.Size([2, 3, 4, 5])
-        assert (stacked_tc.X == 1).all()
-        assert (stacked_tc.y.X == 1).all()
+        if not lazy:
+            assert stacked_tc.X.shape == torch.Size([2, 3, 4, 5])
+            assert stacked_tc.y.X.shape == torch.Size([2, 3, 4, 5])
+
+            assert (stacked_tc.X == 1).all()
+            assert (stacked_tc.y.X == 1).all()
+        else:
+            assert stacked_tc[0].X.shape == torch.Size([3, 4, 5])
+            assert stacked_tc[0].y.X.shape == torch.Size([3, 4, 5])
+            assert stacked_tc[1].X.shape == torch.Size([3, 4, 4])
+            assert stacked_tc[1].y.X.shape == torch.Size([3, 4, 4])
+            assert (stacked_tc[0].X == 1).all()
+            assert (stacked_tc[0].y.X == 1).all()
+
         if lazy_legacy() or lazy:
             assert isinstance(stacked_tc._tensordict, LazyStackedTensorDict)
             assert isinstance(stacked_tc.y._tensordict, LazyStackedTensorDict)