Repository: pytorch/tensordict

Issue:
[Bug] Indexing i:i+1 not handled correctly
```python
from tensordict import tensorclass
import torch

@tensorclass
class MyData:
    x: torch.Tensor
a = MyData(x=torch.ones((3,10)), batch_size=[3])
b = MyData(x=torch.zeros((10)), batch_size=[])
# this works
a[1] = b
# this raises but should not
a[1:2] = b
```


Discussion:
I think this is a broadcasting issue:
```python
a[:] = b
```
should work too
I'm on it

Resolving patch:
diff --git a/tensordict/persistent.py b/tensordict/persistent.py
--- a/tensordict/persistent.py
+++ b/tensordict/persistent.py
@@ -343,16 +343,35 @@ def __getitem__(self, item):
 
     __getitems__ = __getitem__
 
-    def __setitem__(self, item, value):
-        if isinstance(item, str) or (
-            isinstance(item, tuple) and all(isinstance(val, str) for val in item)
+    def __setitem__(self, index, value):
+        if isinstance(index, str) or (
+            isinstance(index, tuple) and all(isinstance(val, str) for val in index)
         ):
-            return self.set(item, value, inplace=True)
+            return self.set(index, value, inplace=True)
 
-        if isinstance(item, list):
+        if isinstance(index, list):
             # convert to tensor
-            item = torch.tensor(item)
-        return self.get_sub_tensordict(item).update(value, inplace=True)
+            index = torch.tensor(index)
+        sub_td = self.get_sub_tensordict(index)
+        err_set_batch_size = None
+        if not isinstance(value, TensorDictBase):
+            value = TensorDict.from_dict(value, batch_size=[])
+            # try to assign the current shape. If that does not work, we can
+            # try to expand
+            try:
+                value.batch_size = sub_td.batch_size
+            except RuntimeError as err0:
+                err_set_batch_size = err0
+        if value.shape != sub_td.shape:
+            try:
+                value = value.expand(sub_td.shape)
+            except RuntimeError as err:
+                if err_set_batch_size is not None:
+                    raise err from err_set_batch_size
+                raise RuntimeError(
+                    f"Cannot broadcast the tensordict {value} to the shape of the indexed persistent tensordict {self}[{index}]."
+                ) from err
+        sub_td.update(value, inplace=True)
 
     def keys(
         self, include_nested: bool = False, leaves_only: bool = False
diff --git a/tensordict/tensordict.py b/tensordict/tensordict.py
--- a/tensordict/tensordict.py
+++ b/tensordict/tensordict.py
@@ -2680,11 +2680,15 @@ def __setitem__(
                     value, batch_size=indexed_bs, device=self.device, _run_checks=False
                 )
             if value.batch_size != indexed_bs:
-                raise RuntimeError(
-                    f"indexed destination TensorDict batch size is {indexed_bs} "
-                    f"(batch_size = {self.batch_size}, index={index}), "
-                    f"which differs from the source batch size {value.batch_size}"
-                )
+                # try to expand
+                try:
+                    value = value.expand(indexed_bs)
+                except RuntimeError as err:
+                    raise RuntimeError(
+                        f"indexed destination TensorDict batch size is {indexed_bs} "
+                        f"(batch_size = {self.batch_size}, index={index}), "
+                        f"which differs from the source batch size {value.batch_size}"
+                    ) from err
 
             keys = set(self.keys())
             if not all(key in keys for key in value.keys()):
@@ -4624,7 +4628,7 @@ def clone(self, recurse: bool = True) -> SubTensorDict:
         return SubTensorDict(source=self._source, idx=self.idx)
 
     def is_contiguous(self) -> bool:
-        return all([value.is_contiguous() for _, value in self.items()])
+        return all(value.is_contiguous() for value in self.values())
 
     def contiguous(self) -> TensorDictBase:
         if self.is_contiguous():
@@ -4647,61 +4651,9 @@ def select(
     def expand(self, *shape: int, inplace: bool = False) -> TensorDictBase:
         if len(shape) == 1 and isinstance(shape[0], Sequence):
             shape = tuple(shape[0])
-
-        idx = self.idx
-        if isinstance(idx, Tensor) and idx.dtype is torch.double:
-            # check that idx is not a mask, otherwise throw an error
-            raise ValueError("Cannot expand a TensorDict masked using SubTensorDict")
-        elif not isinstance(idx, tuple):
-            # create an tuple idx with length equal to this TensorDict's number of dims
-            idx = (idx,) + (slice(None),) * (self._source.ndimension() - 1)
-        elif isinstance(idx, tuple) and len(idx) < self._source.ndimension():
-            # create an tuple idx with length equal to this TensorDict's number of dims
-            idx = idx + (slice(None),) * (self._source.ndimension() - len(idx))
-        # now that idx has the same length as the source's number of dims, we can work with it
-
-        source_shape = self._source.shape
-        num_integer_types = 0
-        for i in idx:
-            if isinstance(i, (int, np.integer)) or (
-                isinstance(i, Tensor) and i.ndimension() == 0
-            ):
-                num_integer_types += 1
-        number_of_extra_dim = len(source_shape) - len(shape) + num_integer_types
-        if number_of_extra_dim > 0:
-            new_source_shape = [shape[i] for i in range(number_of_extra_dim)]
-            shape = shape[len(new_source_shape) :]
-        else:
-            new_source_shape = []
-        new_idx = [slice(None) for _ in range(len(new_source_shape))]
-        for _idx, _s in zip(idx, source_shape):
-            # we're iterating through the source shape and the index
-            # we want to get the new index and the new source shape
-
-            if isinstance(_idx, (int, np.integer)) or (
-                isinstance(_idx, Tensor) and _idx.ndimension() == 0
-            ):
-                # if the index is an integer, do nothing, i.e. keep the index and the shape
-                new_source_shape.append(_s)
-                new_idx.append(_idx)
-            elif _s == 1:
-                # if the source shape at this dim is 1, expand that source dim to the size that is required
-                new_idx.append(slice(None))
-                new_source_shape.append(shape[0])
-                shape = shape[1:]
-            else:
-                # in this case, the source shape must be different than 1. The index is going to be identical.
-                new_idx.append(_idx)
-                new_source_shape.append(shape[0])
-                shape = shape[1:]
-        assert not len(shape)
-        new_source = self._source.expand(*new_source_shape)
-        new_idx = tuple(new_idx)
-        if inplace:
-            self._source = new_source
-            self.idx = new_idx
-            self.batch_size = _getitem_batch_size(new_source_shape, new_idx)
-        return new_source[new_idx]
+        return self.apply(
+            lambda x: x.expand((*shape, *x.shape[self.ndim :])), batch_size=shape
+        )
 
     def is_shared(self) -> bool:
         return self._source.is_shared()


Test patch:
diff --git a/test/test_tensordict.py b/test/test_tensordict.py
--- a/test/test_tensordict.py
+++ b/test/test_tensordict.py
@@ -1471,7 +1471,7 @@ def test_setitem(self, td_name, device, idx):
         td_clone = torch.cat([td_clone, td_clone], 0)
         with pytest.raises(
             RuntimeError,
-            match=r"differs from the source batch size|batch dimension mismatch",
+            match=r"differs from the source batch size|batch dimension mismatch|Cannot broadcast the tensordict",
         ):
             td[idx] = td_clone
 
@@ -2147,6 +2147,37 @@ def test_pop(self, td_name, device):
         ):
             td.pop("z")
 
+    def test_setitem_slice(self, td_name, device):
+        td = getattr(self, td_name)(device)
+        td[:] = td.clone()
+        td[:1] = td[:1].clone().zero_()
+        assert (td[:1] == 0).all()
+        td = getattr(self, td_name)(device)
+        td[:1] = td[:1].to_tensordict().zero_()
+        assert (td[:1] == 0).all()
+
+        # with broadcast
+        td = getattr(self, td_name)(device)
+        td[:1] = td[0].clone().zero_()
+        assert (td[:1] == 0).all()
+        td = getattr(self, td_name)(device)
+        td[:1] = td[0].to_tensordict().zero_()
+        assert (td[:1] == 0).all()
+
+        td = getattr(self, td_name)(device)
+        td[:1, 0] = td[0, 0].clone().zero_()
+        assert (td[:1, 0] == 0).all()
+        td = getattr(self, td_name)(device)
+        td[:1, 0] = td[0, 0].to_tensordict().zero_()
+        assert (td[:1, 0] == 0).all()
+
+        td = getattr(self, td_name)(device)
+        td[:1, :, 0] = td[0, :, 0].clone().zero_()
+        assert (td[:1, :, 0] == 0).all()
+        td = getattr(self, td_name)(device)
+        td[:1, :, 0] = td[0, :, 0].to_tensordict().zero_()
+        assert (td[:1, :, 0] == 0).all()
+
 
 @pytest.mark.parametrize("device", [None, *get_available_devices()])
 @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8])