amend

pytorch · vmoens · Nov 14, 2023 · Jul 21, 2023 · Jul 21, 2023 · Aug 11, 2023
commit d79f4ee4b7a9bdb91ef92e2e72e232f1097453ae
diff --git a/tensordict/memmap_refact.py b/tensordict/memmap_refact.py
@@ -37,7 +37,7 @@ class MemoryMappedTensor(torch.Tensor):
     index: Any
     parent_shape: torch.Size
 
-    def __new__(cls, tensor_or_file, handler=None, dtype=None, shape=None, index=None, device=None):
+    def __new__(cls, tensor_or_file, *, dtype=None, shape=None, index=None, device=None, handler=None):
         if device is not None and torch.device(device).type != "cpu":
             raise ValueError(f"{cls} device must be cpu!")
         if isinstance(tensor_or_file, str):
@@ -63,7 +63,7 @@ def __init__(self, tensor_or_file, handler=None, dtype=None, shape=None, device=
 
     @classmethod
     def from_tensor(
-        cls, tensor, dir=None, prefix=None, filename=None
+        cls, tensor, *, dir=None, filename=None
     ):
         if isinstance(tensor, MemoryMappedTensor):
             if dir is None and (
@@ -112,13 +112,9 @@ def from_tensor(
         out.copy_(tensor)
         return out
 
-    # def __setstate__(self, state: dict[str, Any]) -> None:
-    #     filename = state["filename"]
-    #     handler = state['handler']
-    #     if filename is not None:
-    #         return self.from_filename(filename, state['dtype'], state['shape'])
-    #     else:
-    #         return self.from_handler(handler, state['dtype'], state['shape'])
+    @classmethod
+    def empty_like(cls, tensor, *, filename=None):
+        return cls.from_tensor(torch.zeros((), dtype=tensor.dtype, device=tensor.device).expand_as(tensor), filename=filename)
 
     @classmethod
     def from_filename(cls, filename, dtype, shape, index):
@@ -148,22 +144,22 @@ def from_handler(cls, handler, dtype, shape, index):
         return out
 
     def __reduce__(self):
-        if getattr(self, "handler", None) is not None:
+        if getattr(self, "_handler", None) is not None:
             return type(self).from_handler, (
                 self._handler,
                 self.dtype,
                 self.parent_shape,
                 self.index,
             )
-        elif getattr(self, "filename", None) is not None:
+        elif getattr(self, "_filename", None) is not None:
             return type(self).from_filename, (
                 self._filename,
                 self.dtype,
                 self.parent_shape,
                 self.index,
             )
         else:
-            raise RuntimeError
+            raise RuntimeError("Could not find handler or filename.")
 
     def __getitem__(self, item):
         out = super().__getitem__(item)

diff --git a/tensordict/tensordict.py b/tensordict/tensordict.py
@@ -40,8 +40,8 @@
 import torch
 from functorch import dim as ftdim
 from tensordict._tensordict import _unravel_key_to_tuple
-from tensordict.memmap import memmap_tensor_as_tensor, MemmapTensor
-from tensordict.memmap_refact import MemoryMappedTensor
+from tensordict.memmap import memmap_tensor_as_tensor, MemmapTensor as _MemmapTensor
+from tensordict.memmap_refact import MemoryMappedTensor as MemmapTensor
 from tensordict.utils import (
     _device,
     _dtype,
@@ -132,13 +132,13 @@ def __bool__(self):
 LAZY_TD_HANDLED_FUNCTIONS: dict[Callable, Callable] = {}
 CompatibleType = Union[
     Tensor,
-    MemmapTensor,
+    _MemmapTensor,
 ]  # None? # leaves space for TensorDictBase
 
 if _has_torchrec:
     CompatibleType = Union[
         Tensor,
-        MemmapTensor,
+        _MemmapTensor,
         KeyedJaggedTensor,
     ]
 _STR_MIXED_INDEX_ERROR = "Received a mixed string-non string index. Only string-only or string-free indices are supported."
@@ -1142,8 +1142,8 @@ def _send(self, dst: int, _tag: int = -1, pseudo_rand: bool = False) -> int:
             elif _is_tensor_collection(value.__class__):
                 _tag = value._send(dst, _tag=_tag, pseudo_rand=pseudo_rand)
                 continue
-            elif isinstance(value, MemmapTensor):
-                value = value.as_tensor()
+            # elif isinstance(value, MemmapTensor):
+            #     value = value.as_tensor()
             else:
                 raise NotImplementedError(f"Type {type(value)} is not supported.")
             if not pseudo_rand:
@@ -1181,8 +1181,8 @@ def _recv(self, src: int, _tag: int = -1, pseudo_rand: bool = False) -> int:
             elif _is_tensor_collection(value.__class__):
                 _tag = value._recv(src, _tag=_tag, pseudo_rand=pseudo_rand)
                 continue
-            elif isinstance(value, MemmapTensor):
-                value = value.as_tensor()
+            # elif isinstance(value, MemmapTensor):
+            #     value = value.as_tensor()
             else:
                 raise NotImplementedError(f"Type {type(value)} is not supported.")
             if not pseudo_rand:
@@ -1294,8 +1294,8 @@ def _isend(
                 continue
             elif isinstance(value, Tensor):
                 pass
-            elif isinstance(value, MemmapTensor):
-                value = value.as_tensor()
+            # elif isinstance(value, MemmapTensor):
+            #     value = value.as_tensor()
             else:
                 raise NotImplementedError(f"Type {type(value)} is not supported.")
             if not pseudo_rand:
@@ -1365,8 +1365,8 @@ def _irecv(
                     pseudo_rand=pseudo_rand,
                 )
                 continue
-            elif isinstance(value, MemmapTensor):
-                value = value.as_tensor()
+            # elif isinstance(value, MemmapTensor):
+            #     value = value.as_tensor()
             elif isinstance(value, Tensor):
                 pass
             else:
@@ -1415,8 +1415,8 @@ def _reduce(
                     _future_list=_future_list,
                 )
                 continue
-            elif isinstance(value, MemmapTensor):
-                value = value.as_tensor()
+            # elif isinstance(value, MemmapTensor):
+            #     value = value.as_tensor()
             elif isinstance(value, Tensor):
                 pass
             else:
@@ -2631,8 +2631,8 @@ def to_h5(
             >>>
             >>> from tensordict import TensorDict, MemmapTensor
             >>> td = TensorDict({
-            ...     "a": MemmapTensor(1_000_000),
-            ...     "b": {"c": MemmapTensor(1_000_000, 3)},
+            ...     "a": MemmapTensor.from_tensor(torch.zeros(()).expand(1_000_000)),
+            ...     "b": {"c": MemmapTensor.from_tensor(torch.zeros(()).expand(1_000_000, 3))},
             ... }, [1_000_000])
             >>>
             >>> file = tempfile.NamedTemporaryFile()

diff --git a/test/test_memmap2.py b/test/test_memmap2.py
@@ -382,7 +382,8 @@ def test_handler():
 
 
 def test_memmap_from_memmap():
-    mt2 = MemmapTensor.from_tensor(MemmapTensor(4, 3, 2, 1))
+    mt = MemmapTensor.from_tensor(torch.zeros(()).expand(4, 3, 2, 1))
+    mt2 = MemmapTensor.from_tensor(mt)
     assert mt2.squeeze(-1).shape == torch.Size([4, 3, 2])
 
 
@@ -399,51 +400,43 @@ def test_memmap_cast():
 
 @pytest.fixture
 def dummy_memmap():
-    return MemmapTensor.from_tensor(torch.zeros(10, 11))
+    return MemmapTensor.from_tensor(torch.randn(10, 11))
 
 
-@pytest.mark.parametrize("device", get_available_devices())
 class TestOps:
-    def test_eq(self, device, dummy_memmap):
-        memmap = dummy_memmap.to(device)
+    def test_eq(self, dummy_memmap):
+        memmap = dummy_memmap
         assert (memmap == memmap.clone()).all()
         assert (memmap.clone() == memmap).all()
-        if device.type == "cpu":
-            assert (memmap == memmap.as_tensor()).all()
-            assert (memmap.as_tensor() == memmap).all()
-        else:
-            assert (memmap == memmap._tensor).all()
-            assert (memmap._tensor == memmap).all()
-
-    def test_fill_(self, device, dummy_memmap):
-        memmap = dummy_memmap.to(device)
-        assert (memmap.fill_(1.0) == 1).all()
-
-    def test_copy_(self, device, dummy_memmap):
-        memmap = dummy_memmap.to(device)
-        assert (memmap.copy_(torch.ones(10, 11, device=device)) == 1).all()
-        assert (torch.ones(10, 11, device=device).copy_(memmap) == 1).all()
-
-    def test_or(self, device):
-        memmap = MemmapTensor.from_tensor(torch.ones(10, 11, dtype=torch.bool)).to(
-            device
-        )
+
+    def test_fill_(self, dummy_memmap):
+        memmap = dummy_memmap.fill_(1.0)
+        assert (memmap == 1).all()
+        assert isinstance(memmap, MemmapTensor)
+
+    def test_copy_(self, dummy_memmap):
+        memmap = dummy_memmap.copy_(torch.ones(10, 11))
+        assert (memmap == 1).all()
+        assert isinstance(memmap, MemmapTensor)
+        # check that memmap can be put in a tensor
+        assert (torch.ones(10, 11).copy_(memmap) == 1).all()
+
+    def test_or(self):
+        memmap = MemmapTensor.from_tensor(torch.ones(10, 11, dtype=torch.bool))
         assert (memmap | (~memmap)).all()
 
-    def test_ne(self, device):
-        memmap = MemmapTensor.from_tensor(torch.ones(10, 11, dtype=torch.bool)).to(
-            device
-        )
+    def test_ne(self):
+        memmap = MemmapTensor.from_tensor(torch.ones(10, 11, dtype=torch.bool))
         assert (memmap != ~memmap).all()
 
 
-def test_memmap_del(tmpdir):
-    t = torch.tensor([1])
-    m = MemmapTensor.from_tensor(t, filename=tmpdir / "tensor")
-    # filename = m.filename
-    assert os.path.isfile(tmpdir / "tensor")
-    del m
-    assert not os.path.isfile(tmpdir / "tensor")
+# def test_memmap_del(tmpdir):
+#     t = torch.tensor([1])
+#     m = MemmapTensor.from_tensor(t, filename=tmpdir / "tensor")
+#     # filename = m.filename
+#     assert os.path.isfile(tmpdir / "tensor")
+#     del m
+#     assert not os.path.isfile(tmpdir / "tensor")
 
 # @pytest.mark.parametrize("value", [True, False])
 # def test_memmap_ownership_2pass(value):
@@ -498,40 +491,40 @@ def test_memmap_del(tmpdir):
 #         finally:
 #             p.join()
 #             queue.close()
-@pytest.mark.parametrize(
-    "mode", ["r", "r+", "w+", "c", "readonly", "readwrite", "write", "copyonwrite"]
-)
-def test_mode(mode, tmp_path):
-    mt = MemmapTensor(10, dtype=torch.float32, filename=tmp_path / "test.memmap")
-    mt[:] = torch.ones(10) * 1.5
-    del mt
-
-    if mode in ("r", "readonly"):
-        with pytest.raises(ValueError, match=r"Accepted values for mode are"):
-            MemmapTensor(
-                10, dtype=torch.float32, filename=tmp_path / "test.memmap", mode=mode
-            )
-        return
-    mt = MemmapTensor(
-        10, dtype=torch.float32, filename=tmp_path / "test.memmap", mode=mode
-    )
-    if mode in ("r+", "readwrite", "c", "copyonwrite"):
-        # data in memmap persists
-        assert (mt.as_tensor() == 1.5).all()
-    elif mode in ("w+", "write"):
-        # memmap is initialized to zero
-        assert (mt.as_tensor() == 0).all()
-
-    mt[:] = torch.ones(10) * 2.5
-    assert (mt.as_tensor() == 2.5).all()
-    del mt
-
-    mt2 = MemmapTensor(10, dtype=torch.float32, filename=tmp_path / "test.memmap")
-    if mode in ("c", "copyonwrite"):
-        # tensor was only mutated in memory, not on disk
-        assert (mt2.as_tensor() == 1.5).all()
-    else:
-        assert (mt2.as_tensor() == 2.5).all()
+# @pytest.mark.parametrize(
+#     "mode", ["r", "r+", "w+", "c", "readonly", "readwrite", "write", "copyonwrite"]
+# )
+# def test_mode(mode, tmp_path):
+#     mt = MemmapTensor(10, dtype=torch.float32, filename=tmp_path / "test.memmap")
+#     mt[:] = torch.ones(10) * 1.5
+#     del mt
+#
+#     if mode in ("r", "readonly"):
+#         with pytest.raises(ValueError, match=r"Accepted values for mode are"):
+#             MemmapTensor(
+#                 10, dtype=torch.float32, filename=tmp_path / "test.memmap", mode=mode
+#             )
+#         return
+#     mt = MemmapTensor(
+#         10, dtype=torch.float32, filename=tmp_path / "test.memmap", mode=mode
+#     )
+#     if mode in ("r+", "readwrite", "c", "copyonwrite"):
+#         # data in memmap persists
+#         assert (mt.as_tensor() == 1.5).all()
+#     elif mode in ("w+", "write"):
+#         # memmap is initialized to zero
+#         assert (mt.as_tensor() == 0).all()
+#
+#     mt[:] = torch.ones(10) * 2.5
+#     assert (mt.as_tensor() == 2.5).all()
+#     del mt
+#
+#     mt2 = MemmapTensor(10, dtype=torch.float32, filename=tmp_path / "test.memmap")
+#     if mode in ("c", "copyonwrite"):
+#         # tensor was only mutated in memory, not on disk
+#         assert (mt2.as_tensor() == 1.5).all()
+#     else:
+#         assert (mt2.as_tensor() == 2.5).all()
 
 
 if __name__ == "__main__":