[Feature] More copy() refactors for compile friendliness (#1516)

Author: vmoens

tensordict/_td.py

@@ -926,8 +926,7 @@ def all(self, dim: int | None = None) -> bool | TensorCollection:
 
             names = None
             if self._has_names():
-                names = copy(self.names)
-                names = [name for i, name in enumerate(names) if i != dim]
+                names = [name for i, name in enumerate(self.names) if i != dim]
 
             return TensorDict(
                 source={key: value.all(dim=dim) for key, value in self.items()},
@@ -948,8 +947,7 @@ def any(self, dim: int | None = None) -> bool | TensorCollection:
 
             names = None
             if self._has_names():
-                names = copy(self.names)
-                names = [name for i, name in enumerate(names) if i != dim]
+                names = [name for i, name in enumerate(self.names) if i != dim]
 
             return TensorDict(
                 source={key: value.any(dim=dim) for key, value in self.items()},
@@ -1071,7 +1069,7 @@ def reduction(val):
                 return result
 
             if self._has_names():
-                names = copy(self.names)
+                names = list(self.names)
             else:
                 names = None
             if not call_on_nested:
@@ -1088,11 +1086,10 @@ def reduction(val):
         elif dim is not NO_DEFAULT or keepdim:
             names = None
             if self._has_names():
-                names = copy(self.names)
                 if not keepdim and isinstance(dim, tuple):
-                    names = [name for i, name in enumerate(names) if i not in dim]
+                    names = [name for i, name in enumerate(self.names) if i not in dim]
                 else:
-                    names = [name for i, name in enumerate(names) if i != dim]
+                    names = [name for i, name in enumerate(self.names) if i != dim]
             if dim is not NO_DEFAULT:
                 kwargs["dim"] = dim
             if keepdim is not NO_DEFAULT:
@@ -1745,8 +1742,7 @@ def _unbind(self, dim: int):
         batch_size = torch.Size([s for i, s in enumerate(self.batch_size) if i != dim])
         names = None
         if self._has_names():
-            names = copy(self.names)
-            names = [name for i, name in enumerate(names) if i != dim]
+            names = [name for i, name in enumerate(self.names) if i != dim]
             # We could use any() but dynamo doesn't like generators
             for name in names:
                 if name is not None:
@@ -2072,7 +2068,7 @@ def _permute(tensor):
     def _squeeze(self, dim=None):
         batch_size = self.batch_size
         if dim is None:
-            names = copy(self.names) if self._has_names() else None
+            names = list(self.names) if self._has_names() else None
             if names is not None:
                 batch_size, names = _zip_strict(
                     *[
@@ -2114,7 +2110,7 @@ def _squeeze(tensor):
         batch_size = list(batch_size)
         batch_size.pop(dim)
         batch_size = list(batch_size)
-        names = copy(self.names) if self._has_names() else None
+        names = list(self.names) if self._has_names() else None
         if names:
             names.pop(dim)
 
@@ -2149,7 +2145,7 @@ def _unsqueeze(self, dim: int):
         batch_size.insert(newdim, 1)
         batch_size = torch.Size(batch_size)
 
-        names = copy(self.names) if self._has_names() else None
+        names = list(self.names) if self._has_names() else None
         if names:
             names.insert(newdim, None)
 
@@ -2243,7 +2239,10 @@ def from_dict_instance(
         from tensordict import TensorDict
 
         batch_size_set = torch.Size(()) if batch_size is None else batch_size
-        input_dict = copy(input_dict)
+        if is_compiling():
+            input_dict = type(input_dict)(input_dict)
+        else:
+            input_dict = copy(input_dict)
         for key, value in list(input_dict.items()):
             if isinstance(value, (dict,)):
                 cur_value = self.get(key)

tensordict/base.py

@@ -5560,7 +5560,7 @@ def refine_names(self, *names) -> Self:
 
         """
         # replace ellipsis if any
-        names_copy = copy(names)
+        names_copy = list(names)
         if any(name is Ellipsis for name in names):
             ellipsis_name = [NO_DEFAULT for _ in range(self.ndim - len(names) + 1)]
             names = []
@@ -6053,7 +6053,10 @@ def load_state_dict(
                 return self.update(self_flatten.unflatten_keys("."))
 
         # copy since we'll be using pop
-        state_dict = copy(state_dict)
+        if is_compiling():
+            state_dict = type(state_dict)(state_dict)
+        else:
+            state_dict = copy(state_dict)
         batch_size = state_dict.pop("__batch_size")
         device = state_dict.pop("__device", None)
 
@@ -8959,7 +8962,7 @@ def unflatten(tensor):
             unflatten, batch_size=batch_size, propagate_lock=True, call_on_nested=True
         )
         if self._has_names():
-            names = copy(self.names)
+            names = list(self.names)
             for _ in range(len(unflattened_size) - 1):
                 names.insert(dim, None)
             out.names = names

tensordict/tensorclass.py

@@ -1661,7 +1661,7 @@ def _share_memory_(self):
 
 
 def _load_memmap(cls, prefix: Path, metadata: dict, *, robust_key, **kwargs):
-    non_tensordict = copy(metadata)
+    non_tensordict = dict(metadata)
     del non_tensordict["_type"]
     if os.path.exists(prefix / "other.pickle"):
         with open(prefix / "other.pickle", "rb") as pickle_file:
@@ -2689,7 +2689,7 @@ def _state_dict(
             destination=destination, prefix=prefix, keep_vars=keep_vars, flatten=flatten
         )
     }
-    state_dict["_non_tensordict"] = copy(self._non_tensordict)
+    state_dict["_non_tensordict"] = dict(self._non_tensordict)
     return state_dict
 
 
@@ -3750,9 +3750,7 @@ def _memmap_(
 
         _metadata = {}
         if prefix is not None:
-            _metadata = copy(self._metadata)
-            if _metadata is None:
-                _metadata = {}
+            _metadata = dict(self._metadata) if self._metadata is not None else {}
             _metadata["memmap_prefix"] = prefix
             _metadata["memmaped"] = memmaped
 

test/test_compile.py

@@ -328,6 +328,53 @@ def exclude_keys(td: TensorDict):
         assert "a" in result_c.keys()
         assert "b" not in result_c.keys()
 
+    def test_all_any(self, mode):
+        def call_all(td: TensorDict):
+            return td.all(dim=0)
+
+        def call_any(td: TensorDict):
+            return td.any(dim=0)
+
+        call_all_c = torch.compile(call_all, fullgraph=True, mode=mode)
+        call_any_c = torch.compile(call_any, fullgraph=True, mode=mode)
+
+        data = TensorDict(
+            {"a": torch.tensor([[True, False], [True, True]])},
+            batch_size=[2, 2],
+        )
+
+        result_all = call_all(data)
+        result_all_c = call_all_c(data)
+        assert (result_all["a"] == result_all_c["a"]).all()
+        assert result_all_c.shape == torch.Size([2])
+
+        result_any = call_any(data)
+        result_any_c = call_any_c(data)
+        assert (result_any["a"] == result_any_c["a"]).all()
+        assert result_any_c.shape == torch.Size([2])
+
+    def test_squeeze_unsqueeze(self, mode):
+        def call_squeeze(td: TensorDict):
+            return td.squeeze(0)
+
+        def call_unsqueeze(td: TensorDict):
+            return td.unsqueeze(0)
+
+        call_squeeze_c = torch.compile(call_squeeze, fullgraph=True, mode=mode)
+        call_unsqueeze_c = torch.compile(call_unsqueeze, fullgraph=True, mode=mode)
+
+        data = TensorDict({"a": torch.randn(1, 3)}, batch_size=[1, 3])
+
+        result_squeeze = call_squeeze(data)
+        result_squeeze_c = call_squeeze_c(data)
+        assert result_squeeze.shape == result_squeeze_c.shape
+        assert result_squeeze_c.shape == torch.Size([3])
+
+        result_unsqueeze = call_unsqueeze(result_squeeze)
+        result_unsqueeze_c = call_unsqueeze_c(result_squeeze_c)
+        assert result_unsqueeze.shape == result_unsqueeze_c.shape
+        assert result_unsqueeze_c.shape == torch.Size([1, 3])
+
     def test_names(self, mode):
         import torch._dynamo.exc