[BugFix] Fix improper name setting in __setitem__

ghstack-source-id: 08d0bcf
Pull-Request-resolved: #1313

Author: Vincent Moens

tensordict/_td.py

@@ -1882,6 +1882,7 @@ def rep(leaf):
             batch_size=new_batch_size,
             call_on_nested=True,
             propagate_lock=True,
+            names=self._maybe_names(),
         )
 
     def _repeat(self, *repeats: int) -> TensorDictBase:
@@ -1895,6 +1896,7 @@ def rep(leaf):
             batch_size=new_batch_size,
             call_on_nested=True,
             propagate_lock=True,
+            names=self._maybe_names(),
         )
 
     def _transpose(self, dim0, dim1):

tensordict/base.py

@@ -11712,7 +11712,9 @@ def _validate_value_devicefree(
             if has_names:
                 if value.names[: self.batch_dims] != self.names:
                     # we clone not to corrupt the value
-                    value = value.clone(False).refine_names(*self.names)
+                    value = value.clone(False).refine_names(
+                        *(self.names + value.names[self.batch_dims :])
+                    )
             else:
                 if value._has_names():
                     self.names = value.names[: self.batch_dims]

tensordict/utils.py

@@ -1360,7 +1360,9 @@ def assert_close(
     if not _is_tensor_collection(type(actual)) or not _is_tensor_collection(
         type(expected)
     ):
-        raise TypeError("assert_allclose inputs must be of TensorDict type")
+        raise TypeError(
+            f"assert_allclose inputs must be of TensorDict type, got {type(actual)} and {type(expected)}"
+        )
 
     from tensordict._lazy import LazyStackedTensorDict
 

test/test_tensordict.py

@@ -2232,6 +2232,34 @@ def test_requires_grad(self, device):
         # First stacked tensor has requires_grad == True
         assert list(stacked_td.values())[0].requires_grad is True
 
+    def test_refine_names_setitem_subtd(self):
+        batch_size = 1
+        seq_len = 2
+        n_agents = 3
+        td = TensorDict(
+            {
+                "agents": TensorDict(
+                    {
+                        "obs": torch.zeros((batch_size, seq_len, n_agents, 5)),
+                        "dones": torch.zeros((batch_size, seq_len, n_agents, 1)),
+                    },
+                    batch_size=(batch_size, seq_len, n_agents),
+                    names=[None, "time", "other"],
+                ),
+                "dones": torch.zeros((batch_size, seq_len)),
+            },
+            batch_size=(batch_size, seq_len),
+            names=[None, "time"],
+        )
+        #
+        td["agents"] = td["agents"].repeat_interleave(2, dim=-1)
+        assert len(td["agents"].names) == 3
+        assert td["agents"].names[-1] == "other"
+        td["agents"] = td["agents"].repeat(1, 1, 2)
+        assert td["agents"].names[-1] == "other"
+        td["agents"] = torch.cat((td["agents"], td["agents"]), dim=2)
+        assert td["agents"].names[-1] == "other"
+
     def test_rename_key_nested(self):
         td = TensorDict(a={"b": {"c": 0}})
         td.rename_key_(("a", "b", "c"), ("a", "b"))