Repository: pytorch/tensordict

Issue:
[Feature Request] Operations between TensorDicts
## Motivation

It would be really useful if there would be the possibility to perform operations such as _sum, multiplication_ etc... between two TensorDicts ( key-wise). Is this a feature that you have in mind to introduce in future?

Right now I noticed that is not possible.


## Checklist

- [ X] I have checked that there is no similar issue in the repo (**required**)



Discussion:
Thanks for the question!
Yes that would be cool but what I don't want is to have algebraic operations implemented per se, bc they would only be used by a small subset of tensordicts and they would have undefined behaviours in many cases: this is why we don't have a dtype etc. 
E.g. what would it mean to do `tensordict + 1` if you have a mask in your tensordict? Should you ignore it? Should it be all set to True? Should it become a float?

## What you can already do

Regarding algebraic operations, you can already do `tensordict.apply(lambda x: x+1)` or similar. But I agree with you that we can't do that with 2 tensordicts / tensorclasses.

## What we could implement

I could see this:

```python
import tensordict

td3 = tensordict.apply(lambda x, y: x+y, td1, td2)
```
where td1 and td2 have a matching structure.
We could also do
```python
import tensordict

td4 = tensordict.apply(lambda x, y, z: x+y * z, td1, td2, td3)
```
etc.

That should also cover tensorclass.

Wdyt?
Hi @vmoens , thanks for the quick reply.

I understand you reasoning but I still think that those operation should be available. I have some reasons for that, please correct me if I am missing something here.

_> "E.g. what would it mean to do tensordict + 1 if you have a mask in your tensordict?"_

Well, I would expect that the output would match the exact behaviour of a tensor, if a tensor is masked (lets say a tensor of bools) then it may not make sense to add 1, but it is still doable.

`k = torch.zeros((1,5), dtype=torch.bool)) `
`k + 1`

`//returns:`
`tensor([[1, 1, 1, 1, 1]])`
`
is not reasonable to expect the same from a tensordict? 

We should not make any assumption on the semantics of a tensordict, the content of a tensordict may contain images, a mask,  or random numbers, but at the end are just numbers.
The keys of a dict are just strings and they are representative **only**  for the developer (it is not like a class), that's why it should be up to the developer on which operations apply to its values (like happens with the tensors).

Let me know if this makes sense or if I am missing something.

Thanks

I see your point.
For these use cases you can do `td.apply(lambda x: x+1)` and it returns a tensordict with all the keys updated. Do you think that implementing `__add__` and others would bring a significant value to the lib? 
For instance, we don't implement grad but a nice feature of the lib is that you can store parameters in a tensordict and get all the grads just by doing 

`td.apply(lambda x: x.grad)`

One thing we could do is implement a fallback for whenever a method does not exist, ie if there is not add or matmul calling one of these falls back on calling it over all the tensors and wrapping the results in a tensordict.

Changing topic: what do you think of the apply method with another tensordict?
Should it be a method like 
`def apply_to(self, fun, *other: TensorDict)` or a function like suggested above?
@Daxo32 

We talked about this and I don't think we'll support algebraic operations soon (operations across TDs is a different topic and I'll open a PR with it soon!)

Let me write down the various reasons here:

1. I think apply satisfies these needs already with little effort from developers, and any fallback would actually just call `apply` so the computational advantage of custom native implementations would be at most marginal.
2. A fallback like 
```python
def __getattr__(self, method):
    return lambda *args, **kwargs: self.apply(lambda x: x.method(*args, **kwargs))
```
would work in theory BUT here's the catch. We have not implemented everything yet.
TensorDict is designed to work only across batch dimensions (read: leading dimensions). Now take one func we don't cover yet (say, `Tensor.unique(dim=-1)`). If we do an implicit fallback, `TensorDict.unique(-1)` collect the unique over the last dimension of the "feature" dimensions, not the last dim of the batch. So we'd be buying a ticket for multiple new issues every time a function is not covered, and not because a missing attribute error has been raised but much worse, because the behaviour is not the one intended (ie can cause silent bugs).
3. Along the same line, although unitary ops can sometimes be meaningful, what about `td1 @ td2`? Should we do it across batch dims? I guess many wouldn't expect that. What about feature dims then? Well in that case we'd need to document 
> TensorDict methods work across batch dims only except for `matmul`, `dot` product etc
which would be a bit weird IMO. We would like for TensorDict to stay "intuitive", ie avoid any commitment to behaviours that require long explanations (although this may not be systematically avoidable, as in this case :) )
4. Having unitary ops only would mean that you could potentially have a code like this:
```python
def affine(td1, td2, td3):
    return apply_to(lambda x, y: x @ y, td1, td2) + td3
```
which you can be sure will cause people to come again and ask for matmul to be added to `TensorDict`.

Does that make sense to you?
Again, happy to read your thoughts on this, and see if there's a case for algebraic ops. 
Hello @vmoens , thank you very much ! 
I also thought about this once again and I completely follow your idea and understand you reasoning. The _apply_ method is a good compromise to handle algebraic operations, I am looking forward for the TDs operations! I will follow this repo for sure, it is really interesting.

Thanks :)



Resolving patch:
diff --git a/tensordict/memmap.py b/tensordict/memmap.py
--- a/tensordict/memmap.py
+++ b/tensordict/memmap.py
@@ -31,6 +31,20 @@
 NoneType = type(None)
 EllipsisType = type(Ellipsis)
 
+numpy_to_torch_dtype_dict = {
+    "bool": torch.bool,
+    "uint8": torch.uint8,
+    "int8": torch.int8,
+    "int16": torch.int16,
+    "int32": torch.int32,
+    "int64": torch.int64,
+    "float16": torch.float16,
+    "float32": torch.float32,
+    "float64": torch.float64,
+    "complex64": torch.complex64,
+    "complex128": torch.complex128,
+}
+
 
 def implements_for_memmap(torch_function) -> Callable:
     """Register a torch function override for ScalarTensor."""
@@ -50,7 +64,7 @@ def to_numpy(tensor: Union[torch.Tensor, np.ndarray]) -> np.ndarray:
         return tensor
 
 
-class MemmapTensor(object):
+class MemmapTensor:
     """A torch.tensor interface with a np.memmap array.
 
     A temporary file is created and cleared once the object is out-of-scope.
@@ -88,7 +102,7 @@ class MemmapTensor(object):
 
     Examples:
         >>> x = torch.ones(3,4)
-        >>> x_memmap = MemmapTensor(x)
+        >>> x_memmap = MemmapTensor.from_tensor(x)
         >>> # indexing
         >>> x0 = x_memmap[0]
         >>> x0[:] = 2
@@ -96,7 +110,7 @@ class MemmapTensor(object):
         >>>
         >>> # device
         >>> x = x.to('cuda:0')
-        >>> x_memmap = MemmapTensor(x)
+        >>> x_memmap = MemmapTensor.from_tensor(x)
         >>> assert (x_memmap.clone()).device == torch.device('cuda:0')
         >>>
         >>> # operations
@@ -116,7 +130,6 @@ class MemmapTensor(object):
 
     def __init__(
         self,
-        elem: Union[torch.Tensor, MemmapTensor],
         *size: int,
         device: DEVICE_TYPING = None,
         dtype: torch.dtype = None,
@@ -131,33 +144,71 @@ def __init__(
         self.filename = self.file.name
         self.file.close()  # we close the file for now, but don't delete it
 
-        if isinstance(elem, (torch.Tensor, MemmapTensor, np.ndarray)):
-            if device is not None:
-                raise TypeError(
-                    "device cannot be passed when creating a MemmapTensor from a tensor"
-                )
-            if dtype is not None:
-                raise TypeError(
-                    "dtype cannot be passed when creating a MemmapTensor from a tensor"
-                )
-            self._init_tensor(elem, transfer_ownership)
+        if isinstance(size[0], (torch.Tensor, MemmapTensor, np.ndarray)):
+            raise NotImplementedError(
+                "Creating a Memmap array from a tensor is not permitted anymore. "
+                "Call MemmapTensor.from_tensor(tensor) instead."
+            )
         else:
-            if not isinstance(elem, int) and size:
+            try:
+                shape = (
+                    torch.Size(list(size[0]))
+                    if len(size) == 1 and not isinstance(size[0], int)
+                    else torch.Size(list(size))
+                )
+            except TypeError:
                 raise TypeError(
-                    "Valid init methods for MemmapTensor are: "
-                    "\n- MemmapTensor(tensor, ...)"
-                    "\n- MemmapTensor(size, ...)"
-                    "\n- MemmapTensor(*size, ...)"
+                    f"The *size must be either a single list or tuple of ints, or a sequence of ints. Got {size} instead."
                 )
-            shape = (
-                torch.Size([elem] + list(size))
-                if isinstance(elem, int)
-                else torch.Size(elem)
-            )
             device = device if device is not None else torch.device("cpu")
             dtype = dtype if dtype is not None else torch.get_default_dtype()
-            self._init_shape(shape, device, dtype, transfer_ownership)
-        self._index = None
+            self._init_shape(
+                shape=shape,
+                device=device,
+                dtype=dtype,
+                transfer_ownership=transfer_ownership,
+            )
+        if not hasattr(self, "_index"):
+            self._index = None
+
+    @classmethod
+    def from_tensor(
+        cls,
+        tensor: Union[torch.Tensor, MemmapTensor, np.ndarray],
+        transfer_ownership=False,
+        prefix=None,
+    ) -> MemmapTensor:
+        if isinstance(tensor, MemmapTensor):
+            if transfer_ownership:
+                raise RuntimeError(
+                    "from_tensor(memmap_tensor, transfer_ownership=True) is not permitted, as this method will "
+                    "simply return the original MemmapTensor instance."
+                )
+            return tensor
+        elif isinstance(tensor, np.ndarray):
+            raise TypeError(
+                "Convert input to torch.Tensor before calling MemmapTensor."
+            )
+        if tensor.requires_grad:
+            raise RuntimeError(
+                "MemmapTensor is incompatible with tensor.requires_grad."
+            )
+        device = tensor.device if hasattr(tensor, "device") else torch.device("cpu")
+        dtype = (
+            tensor.dtype
+            if isinstance(tensor, torch.Tensor)
+            else numpy_to_torch_dtype_dict[tensor.dtype.name]
+        )
+        shape = tensor.shape
+        out = cls(
+            shape,
+            device=device,
+            dtype=dtype,
+            prefix=prefix,
+            transfer_ownership=transfer_ownership,
+        )
+        out.copy_(tensor)
+        return out
 
     @staticmethod
     def _create_memmap_with_index(memmap_tensor, index):
@@ -193,49 +244,11 @@ def _init_shape(
         self._numel = prod(shape)
         self.mode = "r+"
         self._has_ownership = True
+        self._had_ownership = True
 
-        self._tensor_dir = torch.zeros(1, device=device, dtype=dtype).__dir__()
+        self._tensor_dir = torch.zeros(0, device=device, dtype=dtype).__dir__()
         self._save_item(shape)
 
-    def _init_tensor(
-        self, elem: Union[torch.Tensor, MemmapTensor], transfer_ownership: bool
-    ):
-        if not isinstance(elem, (torch.Tensor, MemmapTensor)):
-            raise TypeError(
-                "convert input to torch.Tensor before calling MemmapTensor() " "on it."
-            )
-
-        if elem.requires_grad:
-            raise RuntimeError(
-                "MemmapTensor is incompatible with tensor.requires_grad. "
-                "Consider calling tensor.detach() first."
-            )
-
-        self._device = elem.device
-        self._shape = elem.shape
-        self._shape_indexed = None
-        self.transfer_ownership = transfer_ownership
-        self.np_shape = tuple(self._shape)
-        self._dtype = elem.dtype
-        self._tensor_dir = elem.__dir__()
-        self._ndim = elem.ndimension()
-        self._numel = elem.numel()
-        self.mode = "r+"
-        self._has_ownership = True
-        self._had_ownership = True
-        if isinstance(elem, MemmapTensor):
-            prev_filename = elem.filename
-            self._copy_item(prev_filename)
-            # Do we want to raise an error?
-            # if self.memmap_array is elem.memmap_array:
-            #     raise RuntimeError
-        else:
-            if elem.requires_grad:
-                raise Exception(
-                    "memmap is not compatible with gradients, Tensor has requires_grad equals True"
-                )
-            self._save_item(elem)
-
     def _get_memmap_array(self) -> np.memmap:
         if self._memmap_array is None:
             self._memmap_array = np.memmap(
@@ -457,7 +470,12 @@ def set_transfer_ownership(self, value: bool = True) -> MemmapTensor:
     def __deepcopy__(self, memodict=None):
         if memodict is None:
             memodict = {}
-        return MemmapTensor(self)
+        warnings.warn(
+            "calling deepcopy on a memmap tensor involves loading it in memory "
+            "and recreating a memmap tensor from scratch (as no file destination "
+            "can be passed to deepcopy(...)."
+        )
+        return MemmapTensor.from_tensor(self.clone())
 
     def __del__(self) -> None:
         if "_has_ownership" in self.__dir__() and self._has_ownership:
@@ -605,7 +623,7 @@ def to(
             self.device = dest
             return self
         elif isinstance(dest, torch.dtype):
-            return MemmapTensor(self._tensor.to(dest))
+            return MemmapTensor.from_tensor(self._tensor.to(dest))
         else:
             raise NotImplementedError(
                 f"argument dest={dest} to MemmapTensor.to(dest) is not "
diff --git a/tensordict/tensordict.py b/tensordict/tensordict.py
--- a/tensordict/tensordict.py
+++ b/tensordict/tensordict.py
@@ -34,19 +34,17 @@
 import numpy as np
 import torch
 
-from tensordict.utils import _device, _dtype, _get_item, _is_shared, _set_item, _shape
-from torch import Tensor
-from torch.utils._pytree import tree_map
-
-try:
-    from torch.jit._shape_functions import infer_size_impl
-except ImportError:
-    from tensordict.utils import infer_size_impl
-
 from tensordict.memmap import memmap_tensor_as_tensor, MemmapTensor
+
 from tensordict.utils import (
+    _device,
+    _dtype,
+    _get_item,
     _getitem_batch_size,
+    _is_shared,
     _nested_key_type_check,
+    _set_item,
+    _shape,
     _sub_index,
     convert_ellipsis_to_idx,
     DEVICE_TYPING,
@@ -56,6 +54,14 @@
     NESTED_KEY,
     prod,
 )
+from torch import Tensor
+from torch.utils._pytree import tree_map
+
+
+try:
+    from torch.jit._shape_functions import infer_size_impl
+except ImportError:
+    from tensordict.utils import infer_size_impl
 
 # from torch.utils._pytree import _register_pytree_node
 
@@ -594,55 +600,86 @@ def apply_(self, fn: Callable) -> TensorDictBase:
     def apply(
         self,
         fn: Callable,
+        *others: TensorDictBase,
         batch_size: Optional[Sequence[int]] = None,
         inplace: bool = False,
         **constructor_kwargs,
     ) -> TensorDictBase:
         """Applies a callable to all values stored in the tensordict and sets them in a new tensordict.
 
+        The apply method will return an TensorDict instance, regardless of the
+        input type. To keep the same type, one can execute
+
+          >>> out = td.clone(False).update(td.apply(...))
+
         Args:
             fn (Callable): function to be applied to the tensors in the
                 tensordict.
+            *others (TensorDictBase instances, optional): if provided, these
+                tensordicts should have a structure matching the one of the
+                current tensordict. The :obj:`fn` argument should receive as many
+                inputs as the number of tensordicts, including the one where apply is
+                being called.
             batch_size (sequence of int, optional): if provided,
                 the resulting TensorDict will have the desired batch_size.
                 The :obj:`batch_size` argument should match the batch_size after
-                the transformation.
+                the transformation. This is a keyword only argument.
             inplace (bool, optional): if True, changes are made in-place.
-                Default is False.
+                Default is False. This is a keyword only argument.
             **constructor_kwargs: additional keyword arguments to be passed to the
                 TensorDict constructor.
 
         Returns:
             a new tensordict with transformed_in tensors.
 
+        Example:
+            >>> td = TensorDict({"a": -torch.ones(3), "b": {"c": torch.ones(3)}}, batch_size=[3])
+            >>> td_1 = td.apply(lambda x: x+1)
+            >>> assert (td["a"] == 0).all()
+            >>> assert (td["b", "c"] == 2).all()
+            >>> td_2 = td.apply(lambda x, y: x+y, td)
+            >>> assert (td_2["a"] == -2).all()
+            >>> assert (td_2["b", "c"] == 2).all()
         """
-        out = (
-            self
-            if inplace
-            else TensorDict(
+        if inplace:
+            out = self
+        elif batch_size is not None:
+            out = TensorDict(
                 {},
                 batch_size=batch_size,
                 device=self.device,
                 _run_checks=False,
                 **constructor_kwargs,
             )
-            if batch_size is not None
-            else self.clone(recurse=False)
-        )
+        else:
+            out = TensorDict(
+                {},
+                batch_size=self.batch_size,
+                device=self.device,
+                _run_checks=False,
+                **constructor_kwargs,
+            )
+
+        kwargs = {}
+        if not isinstance(self, SubTensorDict):
+            kwargs["_process"] = False
         is_locked = out.is_locked
         if not inplace and is_locked:
             out.unlock()
         for key, item in self.items():
+            _others = [_other[key] for _other in others]
             if isinstance(item, TensorDictBase):
                 item_trsf = item.apply(
-                    fn, inplace=inplace, batch_size=batch_size, **constructor_kwargs
+                    fn,
+                    *_others,
+                    inplace=inplace,
+                    batch_size=batch_size,
+                    **constructor_kwargs,
                 )
             else:
-                item_trsf = fn(item)
+                item_trsf = fn(item, *_others)
             if item_trsf is not None:
-                out.set(
-                    key, item_trsf, inplace=inplace, _run_checks=False, _process=False
-                )
+                out.set(key, item_trsf, inplace=inplace, _run_checks=False, **kwargs)
         if not inplace and is_locked:
             out.lock()
         return out
@@ -1162,15 +1199,19 @@ def unbind(self, dim: int) -> Tuple[TensorDictBase, ...]:
         if dim < 0:
             dim = self.batch_dims + dim
         batch_size = torch.Size([s for i, s in enumerate(self.batch_size) if i != dim])
-        return tuple(
-            self.apply(
-                lambda tensor, idx=_idx: tensor[idx],
-                batch_size=batch_size,
-                _is_shared=self.is_shared(),
-                _is_memmap=self.is_memmap(),
+        out = []
+        for _idx in idx:
+            out.append(
+                self.apply(
+                    lambda tensor, idx=_idx: tensor[idx],
+                    batch_size=batch_size,
+                )
             )
-            for _idx in idx
-        )
+            if self.is_shared():
+                out[-1].share_memory_()
+            elif self.is_memmap():
+                out[-1].memmap_()
+        return tuple(out)
 
     def chunk(self, chunks: int, dim: int = 0) -> Tuple[TensorDictBase, ...]:
         """Splits a tendordict into the specified number of chunks, if possible.
@@ -2658,7 +2699,7 @@ def memmap_(self, prefix=None) -> TensorDictBase:
             if isinstance(value, TensorDictBase):
                 self._tensordict[key] = value.memmap_()
                 continue
-            self._tensordict[key] = MemmapTensor(value, prefix=prefix)
+            self._tensordict[key] = MemmapTensor.from_tensor(value, prefix=prefix)
         self._is_memmap = True
         self.lock()
         return self
@@ -3414,7 +3455,7 @@ def set(
             if self.is_shared() and self.device.type == "cpu":
                 tensor_expand.share_memory_()
             elif self.is_memmap():
-                tensor_expand = MemmapTensor(tensor_expand)
+                tensor_expand = MemmapTensor.from_tensor(tensor_expand)
         parent.set(key, tensor_expand, _run_checks=_run_checks)
         self.set_(key, tensor)
         return self
@@ -4119,7 +4160,10 @@ def clone(self, recurse: bool = True) -> TensorDictBase:
                 *[td.clone() for td in self.tensordicts],
                 stack_dim=self.stack_dim,
             )
-        return LazyStackedTensorDict(*self.tensordicts, stack_dim=self.stack_dim)
+        return LazyStackedTensorDict(
+            *[td.clone(recurse=False) for td in self.tensordicts],
+            stack_dim=self.stack_dim,
+        )
 
     def pin_memory(self) -> TensorDictBase:
         for td in self.tensordicts:


Test patch:
diff --git a/test/test_memmap.py b/test/test_memmap.py
--- a/test/test_memmap.py
+++ b/test/test_memmap.py
@@ -18,23 +18,23 @@
 def test_memmap_type():
     array = np.random.rand(1)
     with pytest.raises(
-        TypeError, match="convert input to torch.Tensor before calling MemmapTensor"
+        TypeError, match="Convert input to torch.Tensor before calling MemmapTensor"
     ):
-        MemmapTensor(array)
+        MemmapTensor.from_tensor(array)
 
 
 def test_grad():
     t = torch.tensor([1.0])
-    MemmapTensor(t)
+    MemmapTensor.from_tensor(t)
     t = t.requires_grad_()
     with pytest.raises(
-        RuntimeError, match="MemmapTensor is incompatible with tensor.requires_grad"
+        RuntimeError, match="MemmapTensor is incompatible with tensor.requires_grad."
     ):
-        MemmapTensor(t)
+        MemmapTensor.from_tensor(t)
     with pytest.raises(
-        RuntimeError, match="MemmapTensor is incompatible with tensor.requires_grad"
+        RuntimeError, match="MemmapTensor is incompatible with tensor.requires_grad."
     ):
-        MemmapTensor(t + 1)
+        MemmapTensor.from_tensor(t + 1)
 
 
 @pytest.mark.parametrize(
@@ -61,7 +61,7 @@ def test_grad():
 def test_memmap_data_type(dtype, shape):
     """Test that MemmapTensor can be created with a given data type and shape."""
     t = torch.tensor([1, 0], dtype=dtype).reshape(shape)
-    m = MemmapTensor(t)
+    m = MemmapTensor.from_tensor(t)
     assert m.dtype == t.dtype
     assert (m == t).all()
     assert m.shape == t.shape
@@ -77,7 +77,7 @@ def test_memmap_data_type(dtype, shape):
 
 def test_memmap_del():
     t = torch.tensor([1])
-    m = MemmapTensor(t)
+    m = MemmapTensor.from_tensor(t)
     filename = m.filename
     assert os.path.isfile(filename)
     del m
@@ -88,7 +88,7 @@ def test_memmap_del():
 @pytest.mark.parametrize("transfer_ownership", [True, False])
 def test_memmap_ownership(transfer_ownership):
     t = torch.tensor([1])
-    m = MemmapTensor(t, transfer_ownership=transfer_ownership)
+    m = MemmapTensor.from_tensor(t, transfer_ownership=transfer_ownership)
     assert not m.file.delete
     with tempfile.NamedTemporaryFile(suffix=".pkl") as tmp:
         pickle.dump(m, tmp)
@@ -118,7 +118,7 @@ def test_memmap_ownership(transfer_ownership):
 @pytest.mark.parametrize("value", [True, False])
 def test_memmap_ownership_2pass(value):
     t = torch.tensor([1])
-    m1 = MemmapTensor(t, transfer_ownership=value)
+    m1 = MemmapTensor.from_tensor(t, transfer_ownership=value)
     with tempfile.NamedTemporaryFile(suffix=".pkl") as tmp2:
         pickle.dump(m1, tmp2)
         m2 = pickle.load(open(tmp2.name, "rb"))
@@ -128,7 +128,7 @@ def test_memmap_ownership_2pass(value):
             assert m1._has_ownership + m2._has_ownership + m3._has_ownership == 1
 
     del m1, m2, m3
-    m1 = MemmapTensor(t, transfer_ownership=value)
+    m1 = MemmapTensor.from_tensor(t, transfer_ownership=value)
     with tempfile.NamedTemporaryFile(suffix=".pkl") as tmp2:
         pickle.dump(m1, tmp2)
         m2 = pickle.load(open(tmp2.name, "rb"))
@@ -138,14 +138,29 @@ def test_memmap_ownership_2pass(value):
             assert m1._has_ownership + m2._has_ownership + m3._has_ownership == 1
 
 
-def test_memmap_new():
+@pytest.mark.parametrize(
+    "index",
+    [
+        None,
+        [
+            0,
+        ],
+    ],
+)
+def test_memmap_new(index):
     t = torch.tensor([1])
-    m1 = MemmapTensor(t)
-    m2 = MemmapTensor(m1)
+    m = MemmapTensor.from_tensor(t)
+    if index is not None:
+        m1 = m[index]
+    else:
+        m1 = m
+    m2 = MemmapTensor.from_tensor(m1)
     assert isinstance(m2, MemmapTensor)
-    assert m2.filename != m1.filename
+    assert m2.filename == m1.filename
     assert m2.filename == m2.file.name
     assert m2.filename == m2.file._closer.name
+    if index is not None:
+        assert m2.contiguous() == t[index]
     m2c = m2.contiguous()
     assert isinstance(m2c, torch.Tensor)
     assert m2c == m1
@@ -158,7 +173,7 @@ def test_memmap_same_device_as_tensor(device):
     Check if device is correct when .to(device) is called.
     """
     t = torch.tensor([1], device=device)
-    m = MemmapTensor(t)
+    m = MemmapTensor.from_tensor(t)
     assert m.device == torch.device(device)
     for other_device in get_available_devices():
         if other_device != device:
@@ -181,7 +196,7 @@ def test_memmap_create_on_same_device(device):
 
 @pytest.mark.parametrize("device", get_available_devices())
 @pytest.mark.parametrize(
-    "value", [torch.zeros([3, 4]), MemmapTensor(torch.zeros([3, 4]))]
+    "value", [torch.zeros([3, 4]), MemmapTensor.from_tensor(torch.zeros([3, 4]))]
 )
 @pytest.mark.parametrize("shape", [[3, 4], [[3, 4]]])
 def test_memmap_zero_value(device, value, shape):
@@ -189,7 +204,7 @@ def test_memmap_zero_value(device, value, shape):
     Test if all entries are zeros when MemmapTensor is created with size.
     """
     value = value.to(device)
-    expected_memmap_tensor = MemmapTensor(value)
+    expected_memmap_tensor = MemmapTensor.from_tensor(value)
     m = MemmapTensor(*shape, device=device)
     assert m.shape == (3, 4)
     assert torch.all(m == expected_memmap_tensor)
@@ -223,7 +238,7 @@ def _recv_and_send(queue, filename, shape):
         del queue
 
     def test_simple_index(self):
-        t = MemmapTensor(torch.zeros(10))
+        t = MemmapTensor.from_tensor(torch.zeros(10))
         # int
         assert isinstance(t[0], MemmapTensor)
         assert t[0].filename == t.filename
@@ -231,7 +246,7 @@ def test_simple_index(self):
         assert t.shape == torch.Size([10])
 
     def test_range_index(self):
-        t = MemmapTensor(torch.zeros(10))
+        t = MemmapTensor.from_tensor(torch.zeros(10))
         # int
         assert isinstance(t[:2], MemmapTensor)
         assert t[:2].filename == t.filename
@@ -239,7 +254,7 @@ def test_range_index(self):
         assert t.shape == torch.Size([10])
 
     def test_double_index(self):
-        t = MemmapTensor(torch.zeros(10))
+        t = MemmapTensor.from_tensor(torch.zeros(10))
         y = t[:2][-1:]
         # int
         assert isinstance(y, MemmapTensor)
@@ -248,14 +263,14 @@ def test_double_index(self):
         assert t.shape == torch.Size([10])
 
     def test_ownership(self):
-        t = MemmapTensor(torch.zeros(10))
+        t = MemmapTensor.from_tensor(torch.zeros(10))
         y = t[:2][-1:]
         del t
         with pytest.raises(FileNotFoundError, match="No such file or directory"):
             y + 0
 
     def test_send_across_procs(self):
-        t = MemmapTensor(torch.zeros(10), transfer_ownership=False)
+        t = MemmapTensor.from_tensor(torch.zeros(10), transfer_ownership=False)
         queue = mp.Queue(1)
         filename = t.filename
         p = mp.Process(
@@ -288,7 +303,7 @@ def test_send_across_procs(self):
             raise e
 
     def test_send_across_procs_index(self):
-        t = MemmapTensor(torch.zeros(10), transfer_ownership=False)
+        t = MemmapTensor.from_tensor(torch.zeros(10), transfer_ownership=False)
         queue = mp.Queue(1)
         filename = t.filename
         p = mp.Process(
@@ -321,18 +336,18 @@ def test_send_across_procs_index(self):
             raise e
 
     def test_iteration(self):
-        t = MemmapTensor(torch.rand(10))
+        t = MemmapTensor.from_tensor(torch.rand(10))
         for i, _t in enumerate(t):
             assert _t == t[i]
 
     def test_iteration_nd(self):
-        t = MemmapTensor(torch.rand(10, 5))
+        t = MemmapTensor.from_tensor(torch.rand(10, 5))
         for i, _t in enumerate(t):
             assert (_t == t[i]).all()
 
     @staticmethod
     def _test_copy_onto_subproc(queue):
-        t = MemmapTensor(torch.rand(10, 5))
+        t = MemmapTensor.from_tensor(torch.rand(10, 5))
         idx = torch.tensor([1, 2])
         queue.put(t[idx], block=True)
         while queue.full():
diff --git a/test/test_tensordict.py b/test/test_tensordict.py
--- a/test/test_tensordict.py
+++ b/test/test_tensordict.py
@@ -739,6 +739,34 @@ def test_zero_(self, td_name, device):
         for k in td.keys():
             assert (td.get(k) == 0).all()
 
+    @pytest.mark.parametrize("inplace", [False, True])
+    def test_apply(self, td_name, device, inplace):
+        td = getattr(self, td_name)(device)
+        td_c = td.to_tensordict()
+        td_1 = td.apply(lambda x: x + 1, inplace=inplace)
+        if inplace:
+            for key in td.keys(True, True):
+                assert (td_c[key] + 1 == td[key]).all()
+                assert (td_1[key] == td[key]).all()
+        else:
+            for key in td.keys(True, True):
+                assert (td_c[key] + 1 != td[key]).any()
+                assert (td_1[key] == td[key] + 1).all()
+
+    @pytest.mark.parametrize("inplace", [False, True])
+    def test_apply_other(self, td_name, device, inplace):
+        td = getattr(self, td_name)(device)
+        td_c = td.to_tensordict()
+        td_1 = td.apply(lambda x, y: x + y, td_c, inplace=inplace)
+        if inplace:
+            for key in td.keys(True, True):
+                assert (td_c[key] * 2 == td[key]).all()
+                assert (td_1[key] == td[key]).all()
+        else:
+            for key in td.keys(True, True):
+                assert (td_c[key] * 2 != td[key]).any()
+                assert (td_1[key] == td[key] * 2).all()
+
     def test_from_empty(self, td_name, device):
         torch.manual_seed(1)
         td = getattr(self, td_name)(device)