[BugFix] Remove _is_memmap and _is_shared from constructor

benchmarks/distributed/distributed_benchmark_test.py

@@ -10,7 +10,7 @@
 import pytest
 import torch
 
-from tensordict import MemoryMappedTensor, TensorDict
+from tensordict import TensorDict
 from torch.distributed import rpc
 
 MAIN_NODE = "Main"
@@ -66,15 +66,12 @@ def exec_distributed_test(rank_node):
             # create a tensordict is 1Gb big, stored on disk, assuming that both nodes have access to /tmp/
             tensordict = TensorDict(
                 {
-                    "memmap": MemoryMappedTensor.empty(
-                        (1000, 640, 640, 3),
-                        dtype=torch.uint8,
-                        filename=tmpdir / "mmap.memmap",
+                    "memmap": torch.empty((), dtype=torch.uint8).expand(
+                        (1000, 640, 640, 3)
                     )
                 },
                 [1000],
-                _is_memmap=True,
-            )
+            ).memmap_(tmpdir, copy_existing=False)
             assert tensordict.is_memmap()
 
             while True:

tensordict/_lazy.py

@@ -2522,8 +2522,6 @@ def exclude(self, *keys: str, inplace: bool = False) -> T:
             batch_size=self.batch_size,
             device=self.device,
             _run_checks=False,
-            _is_memmap=self.is_memmap(),
-            _is_shared=self.is_shared(),
         ).exclude(*keys, inplace=True)
 
     def clone(self, recurse: bool = True) -> T:

tensordict/_td.py

@@ -183,6 +183,8 @@ class TensorDict(TensorDictBase):
     """
 
     _td_dim_names = None
+    _is_shared = False
+    _is_memmap = False
 
     def __init__(
         self,
@@ -191,11 +193,7 @@ def __init__(
         device: DeviceType | None = None,
         names: Sequence[str] | None = None,
         _run_checks: bool = True,
-        _is_shared: bool | None = False,
-        _is_memmap: bool | None = False,
     ) -> None:
-        self._is_shared = _is_shared
-        self._is_memmap = _is_memmap
         if device is not None and isinstance(device, (int, str)):
             device = torch.device(device)
         self._device = device
@@ -211,8 +209,6 @@ def __init__(
                             batch_size=self._batch_size,
                             device=self._device,
                             _run_checks=_run_checks,
-                            _is_shared=_is_shared,
-                            _is_memmap=_is_memmap,
                         )
                     _tensordict[key] = value
             self._td_dim_names = names
@@ -746,8 +742,6 @@ def _convert_to_tensordict(self, dict_value: dict[str, Any]) -> T:
             dict_value,
             batch_size=self.batch_size,
             device=self.device,
-            _is_shared=self._is_shared,
-            _is_memmap=self._is_memmap,
         )
 
     def _index_tensordict(
@@ -788,8 +782,6 @@ def _index_tensordict(
             device=self.device,
             names=names,
             _run_checks=False,
-            _is_shared=self.is_shared(),
-            _is_memmap=self.is_memmap(),
         )
 
     def expand(self, *args, **kwargs) -> T:
@@ -1545,8 +1537,6 @@ def save_metadata(data: TensorDictBase, filepath, metadata=None):
             else TensorDict(
                 {},
                 batch_size=self.batch_size,
-                _is_memmap=True,
-                _is_shared=False,
                 names=self.names if self._has_names() else None,
                 device=torch.device("cpu"),
             )
@@ -1778,8 +1768,6 @@ def clone(self, recurse: bool = True) -> T:
             device=self.device,
             names=copy(self._td_dim_names),
             _run_checks=False,
-            _is_shared=self.is_shared() if not recurse else False,
-            _is_memmap=self.is_memmap() if not recurse else False,
         )
 
     def contiguous(self) -> T:
@@ -1795,8 +1783,6 @@ def empty(self, recurse=False) -> T:
                 source={},
                 names=self._td_dim_names,
                 _run_checks=False,
-                _is_memmap=False,
-                _is_shared=False,
             )
         return super().empty(recurse=recurse)
 
@@ -1837,8 +1823,6 @@ def select(self, *keys: NestedKey, inplace: bool = False, strict: bool = True) -
             # names=self.names if self._has_names() else None,
             names=self._td_dim_names,
             _run_checks=False,
-            _is_memmap=self._is_memmap,
-            _is_shared=self._is_shared,
         )
         if inplace:
             self._tensordict = out._tensordict
@@ -1858,7 +1842,7 @@ def keys(
                 include_nested=include_nested, leaves_only=leaves_only, is_leaf=is_leaf
             )
 
-    # @cache  # noqa: B019
+    @cache  # noqa: B019
     def _nested_keys(
         self,
         include_nested: bool = False,

tensordict/persistent.py

@@ -597,8 +597,6 @@ def save_metadata(data: TensorDictBase, filepath, metadata=None):
             else TensorDict(
                 {},
                 batch_size=self.batch_size,
-                _is_memmap=True,
-                _is_shared=False,
                 names=self.names if self._has_names() else None,
                 device=torch.device("cpu"),
             )
@@ -660,6 +658,8 @@ def _populate(
                 futures.append(
                     executor.submit(save_metadata, dest, prefix / "meta.json", metadata)
                 )
+        dest._is_memmap = True
+        dest.lock_()
         return dest
 
     _load_memmap = TensorDict._load_memmap