pytorch · vmoens · Oct 1, 2024 · Oct 1, 2024
diff --git a/tensordict/base.py b/tensordict/base.py
@@ -2213,6 +2213,74 @@ def reshape(
         """
         ...
 
+    def cat_tensors(
+        self,
+        *keys: NestedKey,
+        out_key: NestedKey,
+        dim: int = 0,
+        keep_entries: bool = False,
+    ) -> T:
+        """Concatenates entries into a new entry and possibly remove the original values.
+
+        Args:
+            keys (sequence of NestedKey): entries to concatenate.
+
+        Keyword Argument:
+            out_key (NestedKey): new key name for the concatenated inputs.
+            keep_entries (bool, optional): if ``False``, entries in ``keys`` will be deleted.
+                Defaults to ``False``.
+            dim (int, optional): the dimension along which the concatenation must occur.
+                Defaults to ``0``.
+
+        Returns: self
+
+        Examples:
+            >>> td = TensorDict(a=torch.zeros(1), b=torch.ones(1))
+            >>> td.cat_tensors("a", "b", out_key="c")
+            >>> assert "a" not in td
+            >>> assert (td["c"] == torch.tensor([0, 1])).all()
+
+        """
+        if keep_entries:
+            entries = [self.get(key) for key in keys]
+        else:
+            entries = [self.pop(key) for key in keys]
+        return self.set(out_key, torch.cat(entries, dim=dim))
+
+    def stack_tensors(
+        self,
+        *keys: NestedKey,
+        out_key: NestedKey,
+        dim: int = 0,
+        keep_entries: bool = False,
+    ) -> T:
+        """Stacks entries into a new entry and possibly remove the original values.
+
+        Args:
+            keys (sequence of NestedKey): entries to stack.
+
+        Keyword Argument:
+            out_key (NestedKey): new key name for the stacked inputs.
+            keep_entries (bool, optional): if ``False``, entries in ``keys`` will be deleted.
+                Defaults to ``False``.
+            dim (int, optional): the dimension along which the stack must occur.
+                Defaults to ``0``.
+
+        Returns: self
+
+        Examples:
+            >>> td = TensorDict(a=torch.zeros(()), b=torch.ones(()))
+            >>> td.stack_tensors("a", "b", out_key="c")
+            >>> assert "a" not in td
+            >>> assert (td["c"] == torch.tensor([0, 1])).all()
+
+        """
+        if keep_entries:
+            entries = [self.get(key) for key in keys]
+        else:
+            entries = [self.pop(key) for key in keys]
+        return self.set(out_key, torch.stack(entries, dim=dim))
+
     @classmethod
     def stack(cls, input, dim=0, *, out=None):
         """Stacks tensordicts into a single tensordict along the given dimension.

diff --git a/test/test_tensordict.py b/test/test_tensordict.py
@@ -3561,6 +3561,25 @@ def test_casts(self, td_name, device):
             for value in tdint.values(True, True, is_leaf=is_leaf)
         )
 
+    @pytest.mark.parametrize("keep_entries", [False, True, None])
+    def test_cat_tensors(self, td_name, device, keep_entries):
+        torch.manual_seed(1)
+        td = getattr(self, td_name)(device)
+        with td.unlock_():
+            a = td.pop("a")
+            td["a"] = a.unsqueeze(-1)
+            td["a_bis"] = td["a"] + 1
+            kwargs = {}
+            if keep_entries is not None:
+                kwargs["keep_entries"] = keep_entries
+            pred_stack = torch.cat([td["a"], td["a_bis"]], -1)
+            td.cat_tensors("a", "a_bis", out_key="cat", dim=-1, **kwargs)
+            assert (td["cat"] == pred_stack).all()
+            if keep_entries:
+                assert "a" in td
+            else:
+                assert "a" not in td
+
     @pytest.mark.parametrize("dim", [0, 1])
     @pytest.mark.parametrize("chunks", [1, 2])
     def test_chunk(self, td_name, device, dim, chunks):
@@ -5984,6 +6003,23 @@ def test_stack_subclasses_on_td(self, td_name, device):
         for key in ("a", "b", "c"):
             assert (stacked_td[key] == td[key]).all()
 
+    @pytest.mark.parametrize("keep_entries", [False, True, None])
+    def test_stack_tensors(self, td_name, device, keep_entries):
+        torch.manual_seed(1)
+        td = getattr(self, td_name)(device)
+        with td.unlock_():
+            td["a_bis"] = td["a"] + 1
+            kwargs = {}
+            if keep_entries is not None:
+                kwargs["keep_entries"] = keep_entries
+            pred_stack = torch.stack([td["a"], td["a_bis"]], -1)
+            td.stack_tensors("a", "a_bis", out_key="stack", dim=-1, **kwargs)
+            assert (td["stack"] == pred_stack).all()
+            if keep_entries:
+                assert "a" in td
+            else:
+                assert "a" not in td
+
     @pytest.mark.filterwarnings("error")
     def test_stack_tds_on_subclass(self, td_name, device):
         torch.manual_seed(1)