better setitem nontensorstack (#1367)

Author: vmoens

tensordict/_lazy.py

@@ -2345,6 +2345,14 @@ def __setitem__(self, index: IndexType, value: T) -> T:
             index = convert_ellipsis_to_idx(index, self.batch_size)
         elif isinstance(index, (list, range)):
             index = torch.as_tensor(index, device=self.device)
+        elif isinstance(index, (type(None), bool)) or (
+            isinstance(index, torch.Tensor)
+            and index.shape == ()
+            and index.dtype == torch.bool
+            and index.all()
+        ):
+            self.unsqueeze(0).update(value)
+            return self
 
         if is_tensor_collection(value) or isinstance(value, dict):
             indexed_bs = _getitem_batch_size(self.batch_size, index)
@@ -2448,6 +2456,13 @@ def __getitem__(self, index: IndexType) -> Any:
                         return leaf.tolist()
                     return leaf.data
                 return leaf
+        if isinstance(index, (type(None), bool)) or (
+            isinstance(index, torch.Tensor)
+            and index.shape == ()
+            and index.dtype == torch.bool
+            and index.all()
+        ):
+            return self.unsqueeze(0)
         split_index = self._split_index(index)
         converted_idx = split_index["index_dict"]
         isinteger = split_index["isinteger"]

tensordict/_td.py

@@ -852,6 +852,18 @@ def __setitem__(
             isinstance(index, tuple) and any(idx is Ellipsis for idx in index)
         ):
             index = convert_ellipsis_to_idx(index, self.batch_size)
+        # Convert index like (True,) or True to (0,) over unsqueezed self
+        if isinstance(index, tuple) and len(index) == 1:
+            index = index[0]
+        if isinstance(index, (bool, type(None))) or (
+            isinstance(index, torch.Tensor)
+            and index.shape == ()
+            and index.dtype == torch.bool
+            and index.all()
+        ):
+            with self.unsqueeze(0) as td_unsqueezed:
+                td_unsqueezed[:] = value
+            return
 
         if isinstance(value, (TensorDictBase, dict)):
             indexed_bs = _getitem_batch_size(self.batch_size, index)
@@ -2531,6 +2543,17 @@ def _set_at_str(self, key, value, idx, *, validated, non_blocking: bool):
                         inplace=False,
                         ignore_lock=True,
                     )
+            # TODO: ultimately, we want to get rid of the above logic
+            # dest_val = dest.maybe_to_stack()
+            # dest_val[idx] = value
+            # if dest_val is not dest:
+            #     self._set_str(
+            #         key,
+            #         dest_val,
+            #         validated=True,
+            #         inplace=False,
+            #         ignore_lock=True,
+            #     )
             return
 
         if isinstance(idx, tuple) and len(idx) and isinstance(idx[0], tuple):

test/test_tensordict.py

@@ -11851,6 +11851,26 @@ def test_new_empty_nontensorstack(self):
         assert isinstance(td.new_empty((4,), empty_lazy=True).get("a"), NonTensorStack)
         assert isinstance(td.new_empty((1,), empty_lazy=True).get("a"), NonTensorStack)
 
+    def test_new_empty_setitem(self):
+        td = TensorDict(
+            a=TensorDict(
+                b=NonTensorStack("a", "b", "c").unsqueeze(-1), batch_size=(3,)
+            ),
+            batch_size=(3,),
+        ).to_lazystack()
+        tdz = td.new_zeros((4,), empty_lazy=True)
+        tdz[torch.tensor([True, True, False, True])] = td
+        assert tdz.get(("a", "b")).tolist() == [["a"], ["b"], ["a"], ["c"]]
+
+    def test_new_empty_setitem_2(self):
+        td = TensorDict(
+            a=TensorDict(b=NonTensorStack("a"), batch_size=(1,)), batch_size=(1,)
+        ).to_lazystack()
+        tdz = td.new_zeros((4,), empty_lazy=True)
+        td["a", "b"] = "new"
+        tdz[torch.tensor([False, False, False, True])] = td
+        assert tdz["a", "b"][-1] == "new"
+
     def test_non_tensor_call(self):
         td0 = TensorDict({"a": 0, "b": 0})
         td1 = TensorDict({"a": 1, "b": 1})