[BugFix] Fix unbind

benchmarks/common/common_ops_test.py

@@ -615,6 +615,21 @@ def main():
     tdc["c", "c", "c"] = c
 
 
+def test_unbind_speed(benchmark):
+    (td,), _ = big_nested_td()
+    benchmark(lambda td: td.unbind(0), td)
+
+
+def test_unbind_speed_stack0(benchmark):
+    (td,), _ = big_nested_stacked_td()
+    benchmark(lambda td: td.unbind(0), td)
+
+
+def test_unbind_speed_stack1(benchmark):
+    (td,), _ = big_nested_stacked_td()
+    benchmark(lambda td: td.unbind(1), td)
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

tensordict/tensordict.py

@@ -2366,10 +2366,6 @@ def unbind(self, dim: int) -> tuple[TensorDictBase, ...]:
         Resulting tensordicts will share the storage of the initial tensordict.
 
         """
-        idx = [
-            ((*tuple(slice(None) for _ in range(dim)), i))
-            for i in range(self.shape[dim])
-        ]
         if dim < 0:
             dim = self.batch_dims + dim
         batch_size = torch.Size([s for i, s in enumerate(self.batch_size) if i != dim])
@@ -2378,16 +2374,18 @@ def unbind(self, dim: int) -> tuple[TensorDictBase, ...]:
             names = copy(names)
             names = [name for i, name in enumerate(names) if i != dim]
         out = []
-        for _idx in idx:
-            out.append(
-                self.apply(
-                    lambda tensor, idx=_idx: tensor[idx],
-                    batch_size=batch_size,
-                )
+        unbind_self_dict = {key: tensor.unbind(dim) for key, tensor in self.items()}
+        for _idx in range(self.batch_size[dim]):
+            td = TensorDict(
+                {key: tensor[_idx] for key, tensor in unbind_self_dict.items()},
+                batch_size=batch_size,
+                _run_checks=False,
+                device=self.device,
+                _is_memmap=False,
+                _is_shared=False,
+                names=names,
             )
-            if names is not None:
-                for item in out:
-                    item.names = names
+            out.append(td)
             if self.is_shared():
                 out[-1].share_memory_()
             elif self.is_memmap():
@@ -3165,6 +3163,7 @@ def unflatten_keys(
                 _run_checks=False,
                 _is_shared=self.is_shared(),
                 _is_memmap=self.is_memmap(),
+                names=self.names,
             )
         else:
             out = self
@@ -3176,7 +3175,9 @@ def unflatten_keys(
                     "Unflattening key(s) in tensordict will override existing unflattened key"
                 )
 
-            tensordict = TensorDict({}, batch_size=self.batch_size, device=self.device)
+            tensordict = TensorDict(
+                {}, batch_size=self.batch_size, device=self.device, names=self.names
+            )
             if key in self.keys():
                 tensordict.update(self[key])
             for old_key, new_key in list_of_keys:

test/test_tensordict.py

@@ -5349,6 +5349,14 @@ def test_from_dict(batch_size, batch_dims, device):
         assert data["d", "g"].batch_size == torch.Size(batch_size)
 
 
+def test_unbind_batchsize():
+    td = TensorDict({"a": TensorDict({"b": torch.zeros(2, 3)}, [2, 3])}, [2])
+    td["a"].batch_size
+    tds = td.unbind(0)
+    assert tds[0].batch_size == torch.Size([])
+    assert tds[0]["a"].batch_size == torch.Size([3])
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)