[BugFix] Add check for split_size in TensorDict.split (#1370)

Co-authored-by: vmoens <vincentmoens@gmail.com>

Author: nkkarpov

tensordict/_td.py

@@ -48,6 +48,8 @@
     _BatchedUninitializedParameter,
     _check_inbuild,
     _clone_value,
+    _create_segments_from_int,
+    _create_segments_from_list,
     _get_item,
     _get_leaf_tensordict,
     _get_shape_from_args,
@@ -1750,71 +1752,28 @@ def split(
         # we must use slices to keep the storage of the tensors
         WRONG_TYPE = "split(): argument 'split_size' must be int or list of ints"
         batch_size = self.batch_size
-        batch_sizes = []
         dim = _maybe_correct_neg_dim(dim, batch_size)
         max_size = batch_size[dim]
         if isinstance(split_size, int):
-            idx0 = 0
-            idx1 = min(max_size, split_size)
-            batch_sizes.append(
-                torch.Size(
-                    tuple(
-                        d if i != dim else idx1 - idx0 for i, d in enumerate(batch_size)
-                    )
-                )
-            )
-            while idx1 < max_size:
-                idx0 = idx1
-                idx1 = min(max_size, idx1 + split_size)
-                batch_sizes.append(
-                    torch.Size(
-                        tuple(
-                            d if i != dim else idx1 - idx0
-                            for i, d in enumerate(batch_size)
-                        )
-                    )
-                )
+            segments = _create_segments_from_int(split_size, max_size)
+            chunks = -(self.batch_size[dim] // -split_size)
+            splits = {k: v.chunk(chunks, dim) for k, v in self.items()}
         elif isinstance(split_size, (list, tuple)):
             if len(split_size) == 0:
                 raise RuntimeError("Insufficient number of elements in split_size.")
-            try:
-                idx0 = 0
-                idx1 = split_size[0]
-                batch_sizes.append(
-                    torch.Size(
-                        tuple(
-                            d if i != dim else idx1 - idx0
-                            for i, d in enumerate(batch_size)
-                        )
-                    )
-                )
-                for idx in split_size[1:]:
-                    idx0 = idx1
-                    idx1 = min(max_size, idx1 + idx)
-                    batch_sizes.append(
-                        torch.Size(
-                            tuple(
-                                d if i != dim else idx1 - idx0
-                                for i, d in enumerate(batch_size)
-                            )
-                        )
-                    )
-            except TypeError:
+            if not all(isinstance(x, int) for x in split_size):
                 raise TypeError(WRONG_TYPE)
-
-            if idx1 < batch_size[dim]:
-                raise RuntimeError(
-                    f"Split method expects split_size to sum exactly to {self.batch_size[dim]} (tensor's size at dimension {dim}), but got split_size={split_size}"
-                )
+            splits = {k: v.split(split_size, dim) for k, v in self.items()}
+            segments = _create_segments_from_list(split_size, max_size)
         else:
             raise TypeError(WRONG_TYPE)
         names = self._maybe_names()
-        # Use chunk instead of split to account for nested tensors if possible
-        if isinstance(split_size, int):
-            chunks = -(self.batch_size[dim] // -split_size)
-            splits = {k: v.chunk(chunks, dim) for k, v in self.items()}
-        else:
-            splits = {k: v.split(split_size, dim) for k, v in self.items()}
+        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))
         ]
@@ -2184,7 +2143,6 @@ def from_dict_instance(
         batch_dims=None,
         names=None,
     ):
-
         if batch_dims is not None and batch_size is not None:
             raise ValueError(
                 "Cannot pass both batch_size and batch_dims to `from_dict`."
@@ -2274,7 +2232,7 @@ def batch_dims(self) -> int:
     @batch_dims.setter
     def batch_dims(self, value: int) -> None:
         raise RuntimeError(
-            f"Setting batch dims on {type(self).__name__} instances is " f"not allowed."
+            f"Setting batch dims on {type(self).__name__} instances is not allowed."
         )
 
     def _has_names(self):
@@ -2763,7 +2721,6 @@ def _memmap_(
         share_non_tensor,
         existsok,
     ) -> T:
-
         if prefix is not None:
             prefix = Path(prefix)
             if not prefix.exists():
@@ -2806,7 +2763,6 @@ def _memmap_(
                     )
                 continue
             else:
-
                 if executor is None:
                     _populate_memmap(
                         dest=dest,

tensordict/utils.py

@@ -8,6 +8,7 @@
 import concurrent.futures
 import functools
 import inspect
+import itertools
 import logging
 
 import math
@@ -3050,3 +3051,36 @@ def _check_is_unflatten(new_shape, old_shape, return_flatten_dim=False):
         #     j = len(new_shape) - j - 1
         return out, (i, j)
     return out
+
+
+def _create_segments_from_int(split_size, max_size):
+    if split_size <= 0:
+        raise RuntimeError(
+            f"split_size must be a positive integer, but got {split_size}."
+        )
+    splits = [
+        (start, min(start + split_size, max_size))
+        for start in range(0, max_size, split_size)
+    ]
+    return splits
+
+
+def _create_segments_from_list(
+    split_size: list[int] | tuple[int],
+    max_size: int,
+):
+    splits = [
+        (start, min(start + size, max_size))
+        for start, size in zip(
+            [0] + list(itertools.accumulate(split_size[:-1])),
+            split_size,
+        )
+    ]
+    total_split_size = sum(split_size)
+    if total_split_size != max_size:
+        raise RuntimeError(
+            f"Split method expects split_size to sum exactly to {max_size}, "
+            f"but got sum({split_size}) = {total_split_size}"
+        )
+
+    return splits

test/test_tensordict.py

@@ -2744,6 +2744,14 @@ def test_split_with_invalid_arguments(self):
             td.split(1, 2)
         with pytest.raises(IndexError, match="Incompatible dim"):
             td.split(1, -3)
+        with pytest.raises(
+            RuntimeError, match="split_size must be a positive integer, but got 0."
+        ):
+            td.split(0, -1)
+        with pytest.raises(
+            RuntimeError, match="split_size must be a positive integer, but got -1."
+        ):
+            td.split(-1, -1)
 
     def test_split_with_negative_dim(self):
         td = TensorDict(