amend

tensordict/_lazy.py

@@ -746,13 +746,7 @@ def _legacy_squeeze(self, dim: int | None = None) -> T:
             stack_dim=stack_dim,
         )
 
-    def unbind(self, dim: int) -> tuple[TensorDictBase, ...]:
-        if dim < 0:
-            dim = self.batch_dims + dim
-        if dim < 0 or dim >= self.ndim:
-            raise ValueError(
-                f"Cannot unbind along dimension {dim} with batch size {self.batch_size}."
-            )
+    def _unbind(self, dim: int) -> tuple[TensorDictBase, ...]:
         if dim == self.stack_dim:
             return tuple(self.tensordicts)
         else:
@@ -2869,7 +2863,7 @@ def _unsqueeze(self, dim):
     all = TensorDict.all
     any = TensorDict.any
     expand = TensorDict.expand
-    unbind = TensorDict.unbind
+    _unbind = TensorDict._unbind
     _get_names_idx = TensorDict._get_names_idx
 
 

tensordict/_td.py

@@ -11,7 +11,6 @@
 import numbers
 import os
 from collections import defaultdict
-from concurrent.futures import ThreadPoolExecutor
 from copy import copy
 from numbers import Number
 from pathlib import Path
@@ -828,9 +827,7 @@ def _expand(tensor):
             _expand, batch_size=shape, call_on_nested=True, names=names
         )
 
-    def unbind(self, dim: int) -> tuple[T, ...]:
-        if dim < 0:
-            dim = self.batch_dims + dim
+    def _unbind(self, dim: int):
         batch_size = torch.Size([s for i, s in enumerate(self.batch_size) if i != dim])
         names = None
         if self._has_names():
@@ -847,7 +844,12 @@ def empty():
         tds = tuple(empty() for _ in range(self.batch_size[dim]))
 
         def unbind(key, val, tds=tds):
-            for td, _val in zip(tds, val.unbind(dim)):
+            unbound = (
+                val.unbind(dim)
+                if not _is_tensor_collection(type(val))
+                else val._unbind(dim)
+            )
+            for td, _val in zip(tds, unbound):
                 td._set_str(key, _val, validated=True, inplace=False)
 
         for key, val in self.items():
@@ -2754,7 +2756,7 @@ def _create_nested_str(self, key):
     reshape = TensorDict.reshape
     split = TensorDict.split
     to_module = TensorDict.to_module
-    unbind = TensorDict.unbind
+    _unbind = TensorDict._unbind
 
     def _view(self, *args, **kwargs):
         raise RuntimeError(

tensordict/base.py

@@ -543,7 +543,6 @@ def expand(self, *args: int | torch.Size) -> T:
         """
         ...
 
-    @abc.abstractmethod
     def unbind(self, dim: int) -> tuple[T, ...]:
         """Returns a tuple of indexed tensordicts, unbound along the indicated dimension.
 
@@ -558,6 +557,21 @@ def unbind(self, dim: int) -> tuple[T, ...]:
             tensor([4, 5, 6, 7])
 
         """
+        batch_dims = self.batch_dims
+        if dim < -batch_dims or dim >= batch_dims:
+            raise RuntimeError(
+                f"the dimension provided ({dim}) is beyond the tensordict dimensions ({self.ndim})."
+            )
+        if dim < 0:
+            dim = batch_dims + dim
+        results = self._unbind(dim)
+        if self._is_memmap or self._is_shared:
+            for result in results:
+                result.lock_()
+        return results
+
+    @abc.abstractmethod
+    def _unbind(self, dim: int) -> tuple[T, ...]:
         ...
 
     def chunk(self, chunks: int, dim: int = 0) -> tuple[TensorDictBase, ...]:

tensordict/nn/params.py

@@ -564,7 +564,7 @@ def chunk(self, chunks: int, dim: int = 0) -> tuple[TensorDictBase, ...]:
         ...
 
     @_fallback
-    def unbind(self, dim: int) -> tuple[TensorDictBase, ...]:
+    def _unbind(self, dim: int) -> tuple[TensorDictBase, ...]:
         ...
 
     @_fallback

tensordict/persistent.py

@@ -1079,7 +1079,7 @@ def _unsqueeze(self, dim):
     reshape = TensorDict.reshape
     split = TensorDict.split
     to_module = TensorDict.to_module
-    unbind = TensorDict.unbind
+    _unbind = TensorDict._unbind
     _get_names_idx = TensorDict._get_names_idx