Update

[ghstack-poisoned]

test/test_distributions.py

@@ -13,6 +13,7 @@
 from _utils_internal import get_default_devices
 from tensordict import TensorDictBase
 from torch import autograd, nn
+from torch.utils._pytree import tree_map
 from torchrl.modules import (
     NormalParamWrapper,
     OneHotCategorical,
@@ -182,7 +183,7 @@ class TestTruncatedNormal:
     @pytest.mark.parametrize("device", get_default_devices())
     def test_truncnormal(self, min, max, vecs, upscale, shape, device):
         torch.manual_seed(0)
-        *vecs, min, max, vecs, upscale = torch.utils._pytree.tree_map(
+        *vecs, min, max, vecs, upscale = tree_map(
             lambda t: torch.as_tensor(t, device=device),
             (*vecs, min, max, vecs, upscale),
         )

torchrl/data/replay_buffers/replay_buffers.py

@@ -31,6 +31,7 @@
 from tensordict.nn.utils import _set_dispatch_td_nn_modules
 from tensordict.utils import expand_as_right, expand_right
 from torch import Tensor
+from torch.utils._pytree import tree_map
 
 from torchrl._utils import _make_ordinal_device, accept_remote_rref_udf_invocation
 from torchrl.data.replay_buffers.samplers import (
@@ -319,9 +320,7 @@ def dim_extend(self, value):
     def _transpose(self, data):
         if is_tensor_collection(data):
             return data.transpose(self.dim_extend, 0)
-        return torch.utils._pytree.tree_map(
-            lambda x: x.transpose(self.dim_extend, 0), data
-        )
+        return tree_map(lambda x: x.transpose(self.dim_extend, 0), data)
 
     def _get_collate_fn(self, collate_fn):
         self._collate_fn = (

torchrl/data/replay_buffers/storages.py

@@ -1367,16 +1367,36 @@ def _collate_list_tensordict(x):
     return out
 
 
+@implement_for("torch", "2.4")
 def _stack_anything(data):
     if is_tensor_collection(data[0]):
         return LazyStackedTensorDict.maybe_dense_stack(data)
-    return torch.utils._pytree.tree_map(
+    return tree_map(
         lambda *x: torch.stack(x),
         *data,
         is_leaf=lambda x: isinstance(x, torch.Tensor) or is_tensor_collection(x),
     )
 
 
+@implement_for("torch", None, "2.4")
+def _stack_anything(data):  # noqa: F811
+    from tensordict import _pytree
+
+    if not _pytree.PYTREE_REGISTERED_TDS:
+        raise RuntimeError(
+            "TensorDict is not registered within PyTree. "
+            "If you see this error, it means tensordicts instances cannot be natively stacked using tree_map. "
+            "To solve this issue, (a) upgrade pytorch to a version > 2.4, or (b) make sure TensorDict is registered in PyTree. "
+            "If this error persists, open an issue on https://github.com/pytorch/rl/issues"
+        )
+    if is_tensor_collection(data[0]):
+        return LazyStackedTensorDict.maybe_dense_stack(data)
+    return tree_map(
+        lambda *x: torch.stack(x),
+        *data,
+    )
+
+
 def _collate_id(x):
     return x