[BugFix] Fix chunk following split fix (#1377)

Author: vmoens

tensordict/_lazy.py

@@ -3740,6 +3740,10 @@ def iter_across_tds():
             for tds in _zip_strict(*tds)
         )
 
+    def chunk(self, chunks: int, dim: int = 0) -> tuple[TensorDictBase, ...]:
+        splits = -(self.batch_size[dim] // -chunks)
+        return self.split(splits, dim)
+
     lock_ = TensorDictBase.lock_
     lock = _renamed_inplace_method(lock_)
 
@@ -4410,6 +4414,10 @@ def _cast_reduction(
             **kwargs,
         )
 
+    def chunk(self, chunks: int, dim: int = 0) -> tuple[TensorDictBase, ...]:
+        splits = -(self.batch_size[dim] // -chunks)
+        return self.split(splits, dim)
+
     __xor__ = TensorDict.__xor__
     __or__ = TensorDict.__or__
     __eq__ = TensorDict.__eq__

tensordict/_td.py

@@ -1790,6 +1790,46 @@ def split(
         )
         return result
 
+    def chunk(self, chunks: int, dim: int = 0) -> tuple[TensorDictBase, ...]:
+        if chunks < 1:
+            raise ValueError(
+                f"chunks must be a strictly positive integer, got {chunks}."
+            )
+        # fall back on split, using upper rounding
+        batch_size = self.batch_size
+        dim = _maybe_correct_neg_dim(dim, batch_size)
+        max_size = batch_size[dim]
+        split_size = -(max_size // -chunks)
+        segments = _create_segments_from_int(split_size, max_size)
+        splits = {k: v.chunk(chunks, dim) for k, v in self.items()}
+        names = self._maybe_names()
+        batch_sizes = [
+            torch.Size(
+                tuple(d if i != dim else end - start for i, d in enumerate(batch_size))
+            )
+            for start, end in segments
+        ]
+        splits = [
+            {k: v[ss] for k, v in splits.items()} for ss in range(len(batch_sizes))
+        ]
+        device = self.device
+        is_shared = self._is_shared
+        is_memmap = self._is_memmap
+        is_locked = self.is_locked
+        result = tuple(
+            self._new_unsafe(
+                source=split,
+                batch_size=bsz,
+                names=names,
+                device=device,
+                lock=is_locked,
+                is_shared=is_shared,
+                is_memmap=is_memmap,
+            )
+            for split, bsz in _zip_strict(splits, batch_sizes)
+        )
+        return result
+
     def masked_select(self, mask: Tensor) -> T:
         d = {}
         mask_expand = mask
@@ -4350,6 +4390,10 @@ def _cast_reduction(
     reshape = TensorDict.reshape
     split = TensorDict.split
 
+    def chunk(self, chunks: int, dim: int = 0) -> tuple[TensorDictBase, ...]:
+        splits = -(self.batch_size[dim] // -chunks)
+        return self.split(splits, dim)
+
     def _view(self, *args, **kwargs):
         raise RuntimeError(
             "Cannot call `view` on a sub-tensordict. Call `reshape` instead."

tensordict/base.py

@@ -3419,6 +3419,7 @@ def unbind(self, dim: int) -> tuple[T, ...]:
     def _unbind(self, dim: int) -> tuple[T, ...]:
         raise NotImplementedError
 
+    @abc.abstractmethod
     def chunk(self, chunks: int, dim: int = 0) -> tuple[TensorDictBase, ...]:
         """Splits a tensordict into the specified number of chunks, if possible.
 
@@ -3443,13 +3444,7 @@ def chunk(self, chunks: int, dim: int = 0) -> tuple[TensorDictBase, ...]:
                      [18, 19]]])
 
         """
-        if chunks < 1:
-            raise ValueError(
-                f"chunks must be a strictly positive integer, got {chunks}."
-            )
-        # fall back on split, using upper rounding
-        split_size = -(self.batch_size[dim] // -chunks)
-        return self.split(split_size, dim=dim)
+        raise NotImplementedError
 
     @overload
     def unsqueeze(self, dim: int) -> T: ...

tensordict/persistent.py

@@ -1429,6 +1429,10 @@ def _unsqueeze(self, dim: int):
             "Cannot call `unsqueeze` on a persistent tensordict. Make it dense before calling this method by calling `to_tensordict`."
         )
 
+    def chunk(self, chunks: int, dim: int = 0) -> tuple[TensorDictBase, ...]:
+        splits = -(self.batch_size[dim] // -chunks)
+        return self.split(splits, dim)
+
     __eq__ = TensorDict.__eq__
     __ne__ = TensorDict.__ne__
     __xor__ = TensorDict.__xor__