Repository: pytorch/tensordict (language: python)

Issue:
[Refactor] Refactor keys, items and values
Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #1058




Discussion / hints:

Fix patch:
diff --git a/tensordict/_lazy.py b/tensordict/_lazy.py
index 88145d71e..a234a4f63 100644
--- a/tensordict/_lazy.py
+++ b/tensordict/_lazy.py
@@ -317,7 +317,7 @@ def _has_exclusive_keys(self):
             return False
 
     @_fails_exclusive_keys
-    def to_dict(self) -> dict[str, Any]: ...
+    def to_dict(self, *, retain_none: bool = True) -> dict[str, Any]: ...
 
     def _reduce_get_metadata(self):
         metadata = {}
@@ -3417,7 +3417,7 @@ def _select(
     ) -> _CustomOpTensorDict:
         if inplace:
             raise RuntimeError("Cannot call select inplace on a lazy tensordict.")
-        return self.to_tensordict()._select(
+        return self.to_tensordict(retain_none=True)._select(
             *keys, inplace=False, strict=strict, set_shared=set_shared
         )
 
diff --git a/tensordict/_td.py b/tensordict/_td.py
index 4387839b5..267667808 100644
--- a/tensordict/_td.py
+++ b/tensordict/_td.py
@@ -4210,7 +4210,7 @@ def _iter():
                 if self.leaves_only:
                     for key in self._keys():
                         target_class = self.tensordict.entry_class(key)
-                        if _is_tensor_collection(target_class):
+                        if not self.is_leaf(target_class):
                             continue
                         yield key
                 else:
@@ -4239,9 +4239,11 @@ def _iter_helper(
                 # For lazy stacks
                 value = value[0]
                 cls = type(value)
+            is_tc = _is_tensor_collection(cls)
+            if self.include_nested and is_tc:
+                if not is_non_tensor(cls):
+                    yield from self._iter_helper(value, prefix=full_key)
             is_leaf = self.is_leaf(cls)
-            if self.include_nested and not is_leaf:
-                yield from self._iter_helper(value, prefix=full_key)
             if not self.leaves_only or is_leaf:
                 yield full_key
 
diff --git a/tensordict/base.py b/tensordict/base.py
index 358cae1b1..fda4445ea 100644
--- a/tensordict/base.py
+++ b/tensordict/base.py
@@ -5570,38 +5570,39 @@ def items(
                 Defaults to ``False``.
 
         """
-        if is_leaf is None:
-            is_leaf = _default_is_leaf
+        if sort:
+            yield from sorted(
+                self.items(
+                    include_nested=include_nested,
+                    leaves_only=leaves_only,
+                    is_leaf=is_leaf,
+                ),
+                key=lambda item: (
+                    item[0] if isinstance(item[0], str) else ".".join(item[0])
+                ),
+            )
+        else:
 
-        def _items():
-            if include_nested and leaves_only:
+            if is_leaf is None:
+                is_leaf = _default_is_leaf
+
+            if include_nested:
                 # check the conditions once only
                 for k in self.keys():
                     val = self._get_str(k, NO_DEFAULT)
-                    if not is_leaf(type(val)):
-                        yield from (
-                            (_unravel_key_to_tuple((k, _key)), _val)
-                            for _key, _val in val.items(
-                                include_nested=include_nested,
-                                leaves_only=leaves_only,
-                                is_leaf=is_leaf,
-                            )
-                        )
-                    else:
+                    cls = type(val)
+                    if not leaves_only or is_leaf(cls):
                         yield k, val
-            elif include_nested:
-                for k in self.keys():
-                    val = self._get_str(k, NO_DEFAULT)
-                    yield k, val
-                    if not is_leaf(type(val)):
-                        yield from (
-                            (_unravel_key_to_tuple((k, _key)), _val)
-                            for _key, _val in val.items(
-                                include_nested=include_nested,
-                                leaves_only=leaves_only,
-                                is_leaf=is_leaf,
+                    if _is_tensor_collection(cls):
+                        if not is_non_tensor(cls):
+                            yield from (
+                                (_unravel_key_to_tuple((k, _key)), _val)
+                                for _key, _val in val.items(
+                                    include_nested=include_nested,
+                                    leaves_only=leaves_only,
+                                    is_leaf=is_leaf,
+                                )
                             )
-                        )
             elif leaves_only:
                 for k in self.keys():
                     val = self._get_str(k, NO_DEFAULT)
@@ -5611,16 +5612,6 @@ def _items():
                 for k in self.keys():
                     yield k, self._get_str(k, NO_DEFAULT)
 
-        if sort:
-            yield from sorted(
-                _items(),
-                key=lambda item: (
-                    item[0] if isinstance(item[0], str) else ".".join(item[0])
-                ),
-            )
-        else:
-            yield from _items()
-
     def non_tensor_items(self, include_nested: bool = False):
         """Returns all non-tensor leaves, maybe recursively."""
         return tuple(
@@ -5657,32 +5648,28 @@ def values(
                 Defaults to ``False``.
 
         """
-        if is_leaf is None:
-            is_leaf = _default_is_leaf
+        if sort:
+            for _, value in self.items(include_nested, leaves_only, is_leaf, sort=sort):
+                yield value
+        else:
+
+            if is_leaf is None:
+                is_leaf = _default_is_leaf
 
-        def _values():
             # check the conditions once only
-            if include_nested and leaves_only:
+            if include_nested:
                 for k in self.keys():
                     val = self._get_str(k, NO_DEFAULT)
-                    if not is_leaf(type(val)):
-                        yield from val.values(
-                            include_nested=include_nested,
-                            leaves_only=leaves_only,
-                            is_leaf=is_leaf,
-                        )
-                    else:
+                    cls = type(val)
+                    if not leaves_only or is_leaf(cls):
                         yield val
-            elif include_nested:
-                for k in self.keys():
-                    val = self._get_str(k, NO_DEFAULT)
-                    yield val
-                    if not is_leaf(type(val)):
-                        yield from val.values(
-                            include_nested=include_nested,
-                            leaves_only=leaves_only,
-                            is_leaf=is_leaf,
-                        )
+                    if include_nested and _is_tensor_collection(cls):
+                        if not is_non_tensor(cls):
+                            yield from val.values(
+                                include_nested=include_nested,
+                                leaves_only=leaves_only,
+                                is_leaf=is_leaf,
+                            )
             elif leaves_only:
                 for k in self.keys(sort=sort):
                     val = self._get_str(k, NO_DEFAULT)
@@ -5692,12 +5679,6 @@ def _values():
                 for k in self.keys(sort=sort):
                     yield self._get_str(k, NO_DEFAULT)
 
-        if not sort or not include_nested:
-            yield from _values()
-        else:
-            for _, value in self.items(include_nested, leaves_only, is_leaf, sort=sort):
-                yield value
-
     @cache  # noqa: B019
     def _values_list(
         self,
@@ -9350,9 +9331,16 @@ def _maybe_set_shared_attributes(self, result, lock=False):
             if lock:
                 result.lock_()
 
-    def to_tensordict(self) -> T:
+    def to_tensordict(self, *, retain_none: bool | None = None) -> T:
         """Returns a regular TensorDict instance from the TensorDictBase.
 
+        Args:
+            retain_none (bool): if ``True``, the ``None`` values from tensorclass instances
+                will be written in the tensordict.
+                Otherwise they will be discarded. Default: ``True``.
+
+                .. note:: from v0.8, the default value will be switched to ``False``.
+
         Returns:
             a new TensorDict object containing the same values.
 
@@ -9364,7 +9352,11 @@ def to_tensordict(self) -> T:
                 key: (
                     value.clone()
                     if not _is_tensor_collection(type(value))
-                    else value if is_non_tensor(value) else value.to_tensordict()
+                    else (
+                        value
+                        if is_non_tensor(value)
+                        else value.to_tensordict(retain_none=retain_none)
+                    )
                 )
                 for key, value in self.items(is_leaf=_is_leaf_nontensor)
             },
@@ -9467,12 +9459,27 @@ def as_tensor(tensor):
 
         return self._fast_apply(as_tensor, propagate_lock=True)
 
-    def to_dict(self) -> dict[str, Any]:
-        """Returns a dictionary with key-value pairs matching those of the tensordict."""
-        return {
-            key: value.to_dict() if _is_tensor_collection(type(value)) else value
-            for key, value in self.items()
-        }
+    def to_dict(self, *, retain_none: bool = True) -> dict[str, Any]:
+        """Returns a dictionary with key-value pairs matching those of the tensordict.
+
+        Args:
+            retain_none (bool): if ``True``, the ``None`` values from tensorclass instances
+                will be written in the dictionary.
+                Otherwise, they will be discarded. Default: ``True``.
+
+        """
+        result = {}
+        for key, value in self.items():
+            if _is_tensor_collection(type(value)):
+                if (
+                    not retain_none
+                    and _is_non_tensor(type(value))
+                    and value.data is None
+                ):
+                    continue
+                value = value.to_dict(retain_none=retain_none)
+            result[key] = value
+        return result
 
     def numpy(self):
         """Converts a tensordict to a (possibly nested) dictionary of numpy arrays.
@@ -9500,7 +9507,7 @@ def numpy(self):
             {'a': {'b': array(0., dtype=float32), 'c': 'a string!'}}
 
         """
-        as_dict = self.to_dict()
+        as_dict = self.to_dict(retain_none=False)
 
         def to_numpy(x):
             if isinstance(x, torch.Tensor):
@@ -9541,7 +9548,7 @@ def dict_to_namedtuple(dictionary):
             )
             return cls(**dictionary)
 
-        return dict_to_namedtuple(self.to_dict())
+        return dict_to_namedtuple(self.to_dict(retain_none=False))
 
     @classmethod
     def from_namedtuple(cls, named_tuple, *, auto_batch_size: bool = False):
diff --git a/tensordict/nn/params.py b/tensordict/nn/params.py
index 0116825a0..e4575ef7a 100644
--- a/tensordict/nn/params.py
+++ b/tensordict/nn/params.py
@@ -642,7 +642,7 @@ def from_dict(cls, *args, **kwargs):
         return TensorDictParams(td)
 
     @_fallback
-    def to_tensordict(self): ...
+    def to_tensordict(self, *, retain_none: bool | None = None): ...
 
     @_fallback
     def to_h5(
diff --git a/tensordict/tensorclass.py b/tensordict/tensorclass.py
index 8906eefd4..80475631f 100644
--- a/tensordict/tensorclass.py
+++ b/tensordict/tensorclass.py
@@ -857,7 +857,7 @@ def get_parent_locals(cls, localns=localns):
         cls._type_hints = None
 
 
-def _from_tensordict(cls, tensordict, non_tensordict=None):  # noqa: D417
+def _from_tensordict(cls, tensordict, non_tensordict=None, safe=True):  # noqa: D417
     """Tensor class wrapper to instantiate a new tensor class object.
 
     Args:
@@ -865,7 +865,7 @@ def _from_tensordict(cls, tensordict, non_tensordict=None):  # noqa: D417
         non_tensordict (dict): Dictionary with non-tensor and nested tensor class objects
 
     """
-    if not isinstance(tensordict, TensorDictBase):
+    if safe and not isinstance(tensordict, TensorDictBase):
         raise RuntimeError(
             f"Expected a TensorDictBase instance but got {type(tensordict)}"
         )
@@ -890,10 +890,11 @@ def _from_tensordict(cls, tensordict, non_tensordict=None):  # noqa: D417
         exp_keys = set(cls.__expected_keys__)
         if non_tensordict is not None:
             nontensor_keys = set(non_tensordict.keys())
+            total_keys = tensor_keys.union(nontensor_keys)
         else:
             nontensor_keys = set()
             non_tensordict = {}
-        total_keys = tensor_keys.union(nontensor_keys)
+            total_keys = tensor_keys
     for key in nontensor_keys:
         if key not in tensor_keys:
             continue
@@ -917,11 +918,12 @@ def _from_tensordict(cls, tensordict, non_tensordict=None):  # noqa: D417
         # empty tensordict and writing values to it. we can skip this because we already
         # have a tensordict to use as the underlying tensordict
         tc = cls.__new__(cls)
-        tc.__dict__["_tensordict"] = tensordict
-        tc.__dict__["_non_tensordict"] = non_tensordict
+        tc.__dict__.update(
+            {"_tensordict": tensordict, "_non_tensordict": non_tensordict}
+        )
         # since we aren't calling the dataclass init method, we need to manually check
         # whether a __post_init__ method has been defined and invoke it if so
-        if hasattr(tc, "__post_init__"):
+        if hasattr(cls, "__post_init__"):
             tc.__post_init__()
         return tc
     else:
@@ -1143,6 +1145,23 @@ def wrapper(self, key: str, value: Any) -> None:  # noqa: D417
 
 
 def _wrap_td_method(funcname, *, copy_non_tensor=False, no_wrap=False):
+    def deliver_result(self, result, kwargs):
+        if result is None:
+            return
+        if isinstance(result, TensorDictBase) and kwargs.get("out") is not result:
+            if not is_dynamo_compiling():
+                non_tensordict = super(type(self), self).__getattribute__(
+                    "_non_tensordict"
+                )
+            else:
+                non_tensordict = self._non_tensordict
+            non_tensordict = dict(non_tensordict)
+            if copy_non_tensor and non_tensordict:
+                # use tree_map to copy
+                non_tensordict = tree_map(lambda x: x, non_tensordict)
+            return self._from_tensordict(result, non_tensordict, safe=False)
+        return result
+
     def wrapped_func(self, *args, **kwargs):
         if not is_dynamo_compiling():
             td = super(type(self), self).__getattribute__("_tensordict")
@@ -1154,34 +1173,12 @@ def wrapped_func(self, *args, **kwargs):
         if no_wrap:
             return result
 
-        def check_out(kwargs, result):
-            out = kwargs.get("out")
-            if out is result:
-                # No need to transform output
-                return True
-            return False
-
         if result is td:
             return self
 
-        def deliver_result(result):
-            if isinstance(result, TensorDictBase) and not check_out(kwargs, result):
-                if not is_dynamo_compiling():
-                    non_tensordict = super(type(self), self).__getattribute__(
-                        "_non_tensordict"
-                    )
-                else:
-                    non_tensordict = self._non_tensordict
-                non_tensordict = dict(non_tensordict)
-                if copy_non_tensor:
-                    # use tree_map to copy
-                    non_tensordict = tree_map(lambda x: x, non_tensordict)
-                return self._from_tensordict(result, non_tensordict)
-            return result
-
         if isinstance(result, tuple):
-            return tuple(deliver_result(r) for r in result)
-        return deliver_result(result)
+            return tuple(deliver_result(self, r, kwargs) for r in result)
+        return deliver_result(self, result, kwargs)
 
     return wrapped_func
 
@@ -1428,10 +1425,16 @@ def _len(self) -> int:
     return len(self._tensordict)
 
 
-def _to_dict(self) -> dict:
-    td_dict = self._tensordict.to_dict()
+def _to_dict(self, *, retain_none: bool = True) -> dict:
+    td_dict = self._tensordict.to_dict(retain_none=retain_none)
     if self._non_tensordict:
-        td_dict.update(self._non_tensordict)
+        if retain_none:
+            td_dict.update(self._non_tensordict)
+        else:
+            td_dict.update(
+                {k: v for k, v in self._non_tensordict.items() if v is not None}
+            )
+
     return td_dict
 
 
@@ -1499,20 +1502,36 @@ def _from_dict_instance(
     return out
 
 
-def _to_tensordict(self) -> TensorDict:
+def _to_tensordict(self, *, retain_none: bool | None = None) -> TensorDict:
     """Convert the tensorclass into a regular TensorDict.
 
     Makes a copy of all entries. Memmap and shared memory tensors are converted to
     regular tensors.
 
+    Args:
+        retain_none (bool): if ``True``, the ``None`` values will be written in the
+            tensordict. Otherwise they will be discrarded. Default: ``True``.
+
+            .. note:: from v0.8, the default value will be switched to ``False``.
+
     Returns:
         A new TensorDict object containing the same values as the tensorclass.
 
     """
-    td = self._tensordict.to_tensordict()
+    td = self._tensordict.to_tensordict(retain_none=retain_none)
     for key, val in self._non_tensordict.items():
-        # if val is None:
-        #     continue
+        if val is None:
+            if retain_none is None:
+                retain_none = True
+                warnings.warn(
+                    "retain_none was not specified and a None value was encountered in the tensorclass. "
+                    "As of now, the None will be written in the tensordict but this default behaviour will change"
+                    "in v0.8. To disable this warning, specify the value of retain_none."
+                )
+            if retain_none:
+                pass
+            else:
+                continue
         td.set_non_tensor(key, val)
     return td
 
@@ -2674,11 +2693,11 @@ def _apply_nest(self, *args, out=None, **kwargs):
             names=kwargs.get("names"),
         )
 
-    def to_dict(self):
+    def to_dict(self, *, retain_none: bool = True):
         # override to_dict to return just the data
         return self.data
 
-    def to_tensordict(self):
+    def to_tensordict(self, *, retain_none: bool | None = None):
         return self
 
     @classmethod
@@ -2952,10 +2971,10 @@ def lazy_stack(
             )
         return result
 
-    def to_dict(self) -> dict[str, Any]:
+    def to_dict(self, *, retain_none: bool = True) -> dict[str, Any]:
         return self.tolist()
 
-    def to_tensordict(self):
+    def to_tensordict(self, *, retain_none: bool | None = None):
         return self
 
     def _memmap_(
diff --git a/tensordict/utils.py b/tensordict/utils.py
index 280b224a0..fc40e78d9 100644
--- a/tensordict/utils.py
+++ b/tensordict/utils.py
@@ -1511,16 +1511,12 @@ def assert_close(
         return True
 
     try:
-        set1 = set(
-            actual.keys(is_leaf=lambda x: not is_non_tensor(x), leaves_only=True)
-        )
-        set2 = set(
-            expected.keys(is_leaf=lambda x: not is_non_tensor(x), leaves_only=True)
-        )
-    except ValueError:
-        # Persistent tensordicts do not work with is_leaf
         set1 = set(actual.keys())
         set2 = set(expected.keys())
+    except ValueError:
+        # Persistent tensordicts do not work with is_leaf
+        set1 = set(actual.keys(is_leaf=lambda cls: issubclass(cls, torch.Tensor)))
+        set2 = set(expected.keys(is_leaf=lambda cls: issubclass(cls, torch.Tensor)))
     if not (len(set1.difference(set2)) == 0 and len(set2) == len(set1)):
         raise KeyError(
             "actual and expected tensordict keys mismatch, "


Test patch:
diff --git a/test/test_tensorclass.py b/test/test_tensorclass.py
index b33a4a187..df6d41f68 100644
--- a/test/test_tensorclass.py
+++ b/test/test_tensorclass.py
@@ -1006,7 +1006,10 @@ def test_pickle(
             with open(tempdir / "test.pkl", "rb") as f:
                 data2 = pickle.load(f)
 
-        assert_allclose_td(data.to_tensordict(), data2.to_tensordict())
+        assert_allclose_td(
+            data.to_tensordict(retain_none=False),
+            data2.to_tensordict(retain_none=False),
+        )
         assert isinstance(data2, MyData)
         assert data2.z == data.z
 
@@ -1524,14 +1527,14 @@ def get_data(cls, shift):
         assert (data0.y.X == 1).all()
         assert data0.y.z == "test_tensorclass1"
         data0 = MyDataNested.get_data(0)
-        data0.update(data1.to_dict())
+        data0.update(data1.to_dict(retain_none=False))
         assert (data0.X == 1).all()
         assert data0.z == "test_tensorclass1", data0.z
         assert (data0.y.X == 1).all()
         assert data0.y.z == "test_tensorclass1"
 
         data0 = MyDataNested.get_data(0)
-        data0.update(data1.to_tensordict())
+        data0.update(data1.to_tensordict(retain_none=False))
         assert (data0.X == 1).all()
         assert data0.z == "test_tensorclass1"
         assert (data0.y.X == 1).all()
@@ -1745,7 +1748,7 @@ class MyClass:
             batch_size=[3, 4],
         )
 
-        ctd = c.to_tensordict()
+        ctd = c.to_tensordict(retain_none=False)
         assert isinstance(ctd, TensorDictBase)
         assert "x" in ctd.keys()
         assert "z" in ctd.keys()
diff --git a/test/test_tensordict.py b/test/test_tensordict.py
index 0f1f65b5d..08b0c2384 100644
--- a/test/test_tensordict.py
+++ b/test/test_tensordict.py
@@ -228,7 +228,7 @@ def test_batchsize_reset(self):
             ),
         ):
             td_stack.batch_size = [2]
-        td_stack.to_tensordict().batch_size = [2]
+        td_stack.to_tensordict(retain_none=True).batch_size = [2]
 
         td = TensorDict({"a": torch.randn(3, 4)}, [3, 4])
         subtd = td._get_sub_tensordict((slice(None), torch.tensor([1, 2])))
@@ -239,7 +239,7 @@ def test_batchsize_reset(self):
             ),
         ):
             subtd.batch_size = [3, 2]
-        subtd.to_tensordict().batch_size = [3, 2]
+        subtd.to_tensordict(retain_none=True).batch_size = [3, 2]
 
         td = TensorDict({"a": torch.randn(3, 4)}, [3, 4])
         with set_lazy_legacy(True):
@@ -251,7 +251,7 @@ def test_batchsize_reset(self):
                 ),
             ):
                 td_u.batch_size = [1]
-            td_u.to_tensordict().batch_size = [1]
+            td_u.to_tensordict(retain_none=True).batch_size = [1]
 
     @pytest.mark.parametrize("count_duplicates", [False, True])
     def test_bytes(self, count_duplicates, device_fixture):
@@ -828,7 +828,7 @@ def test_filling_empty_tensordict(self, device, td_type, update):
         for i in range(16):
             other_td = TensorDict({"a": torch.randn(10), "b": torch.ones(1)}, [])
             if td_type == "unsqueeze":
-                other_td = other_td.unsqueeze(-1).to_tensordict()
+                other_td = other_td.unsqueeze(-1).to_tensordict(retain_none=True)
             if update:
                 subtd = td._get_sub_tensordict(i)
                 subtd.update(other_td, inplace=True, non_blocking=False)
@@ -3234,7 +3234,7 @@ def test_add_batch_dim_cache(self, td_name, device, nested):
     @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_c = td.to_tensordict(retain_none=True)
         if inplace and td_name == "td_params":
             with pytest.raises(ValueError, match="Failed to update"):
                 td.apply(lambda x: x + 1, inplace=inplace)
@@ -3254,7 +3254,7 @@ def test_apply_default(self, td_name, device, inplace):
         if td_name in ("td_h5",):
             pytest.skip("Cannot test assignment in persistent tensordict.")
         td = getattr(self, td_name)(device)
-        td_c = td.to_tensordict()
+        td_c = td.to_tensordict(retain_none=True)
         if td_name in ("td_params",):
             td.data.zero_()
         else:
@@ -3368,7 +3368,7 @@ def func(x):
     @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_c = td.to_tensordict(retain_none=True)
         if inplace and td_name == "td_params":
             td_set = td.data
         else:
@@ -3386,12 +3386,13 @@ def test_apply_other(self, td_name, device, inplace):
     def test_apply_out(self, td_name, device):
         td = getattr(self, td_name)(device)
         if not isinstance(td, LazyStackedTensorDict):
-            td_c = td.to_tensordict()
+            td_c = td.to_tensordict(retain_none=True)
         else:
             td_c = td.clone()
         td.apply(lambda x: x + 1, out=td_c)
         assert_allclose_td(
-            td.filter_non_tensor_data().data + 1, td_c.filter_non_tensor_data()
+            td.filter_non_tensor_data().data + 1,
+            td_c.filter_non_tensor_data().filter_empty_(),
         )
 
     def test_as_tensor(self, td_name, device):
@@ -3437,7 +3438,7 @@ def test_auto_batch_size_(self, td_name, device):
     def test_broadcast(self, td_name, device):
         torch.manual_seed(1)
         td = getattr(self, td_name)(device)
-        sub_td = td[:, :2].to_tensordict()
+        sub_td = td[:, :2].to_tensordict(retain_none=True)
         sub_td.zero_()
         sub_dict = sub_td.to_dict()
         if td_name == "td_params":
@@ -3849,15 +3850,15 @@ def test_entry_type(self, td_name, device):
     def test_equal(self, td_name, device):
         torch.manual_seed(1)
         td = getattr(self, td_name)(device)
-        assert (td == td.to_tensordict()).all()
-        td0 = td.to_tensordict().zero_()
+        assert (td == td.to_tensordict(retain_none=True)).all()
+        td0 = td.to_tensordict(retain_none=True).zero_()
         assert (td != td0).any()
 
     def test_equal_dict(self, td_name, device):
         torch.manual_seed(1)
         td = getattr(self, td_name)(device)
         assert (td == td.to_dict()).all()
-        td0 = td.to_tensordict().zero_().to_dict()
+        td0 = td.to_tensordict(retain_none=True).zero_().to_dict()
         assert (td != td0).any()
 
     def test_equal_float(self, td_name, device):
@@ -3886,7 +3887,7 @@ def test_equal_int(self, td_name, device):
             td_set = td
         td_set.zero_()
         assert (td == 0).all()
-        td0 = td.to_tensordict().zero_()
+        td0 = td.to_tensordict(retain_none=True).zero_()
         assert (td0 != 1).all()
 
     def test_equal_other(self, td_name, device):
@@ -3903,7 +3904,7 @@ def test_equal_tensor(self, td_name, device):
             td_set = td
         td_set.zero_()
         assert (td == torch.zeros([], dtype=torch.int, device=device)).all()
-        td0 = td.to_tensordict().zero_()
+        td0 = td.to_tensordict(retain_none=True).zero_()
         assert (td0 != torch.ones([], dtype=torch.int, device=device)).all()
 
     def test_exclude(self, td_name, device):
@@ -4092,7 +4093,7 @@ def test_flatten_unflatten(self, td_name, device):
         shape = td.shape[:3]
         td_flat = td.flatten(0, 2)
         td_unflat = td_flat.unflatten(0, shape)
-        assert (td.to_tensordict() == td_unflat).all()
+        assert (td.to_tensordict(retain_none=True) == td_unflat).all()
         assert td.batch_size == td_unflat.batch_size
 
     @pytest.mark.parametrize("start_dim", [0, 1, -2, -3])
@@ -4129,7 +4130,7 @@ def test_flatten_unflatten_bis(self, td_name, device):
         shape = td.shape[1:4]
         td_flat = td.flatten(1, 3)
         td_unflat = td_flat.unflatten(1, shape)
-        assert (td.to_tensordict() == td_unflat).all()
+        assert (td.to_tensordict(retain_none=True) == td_unflat).all()
         assert td.batch_size == td_unflat.batch_size
 
     def test_from_empty(self, td_name, device):
@@ -4183,7 +4184,7 @@ def test_getitem_ellipsis(self, td_name, device, actual_index, expected_index):
 
         actual_td = td[actual_index]
         expected_td = td[expected_index]
-        other_expected_td = td.to_tensordict()[expected_index]
+        other_expected_td = td.to_tensordict(retain_none=False)[expected_index]
         assert expected_td.shape == _getitem_batch_size(
             td.batch_size, convert_ellipsis_to_idx(actual_index, td.batch_size)
         )
@@ -4245,9 +4246,12 @@ def test_index_none(self, td_name, device, idx):
         # Fixed by 451
         # if td_name == "td_h5":
         #     with pytest.raises(TypeError, match="can't process None"):
-        #         assert (tdnone.to_tensordict() == td.to_tensordict()[idx]).all()
+        #         assert (tdnone.to_tensordict(retain_none=True) == td.to_tensordict(retain_none=True)[idx]).all()
         #     return
-        assert (tdnone.to_tensordict() == td.to_tensordict()[idx]).all()
+        assert (
+            tdnone.to_tensordict(retain_none=True)
+            == td.to_tensordict(retain_none=True)[idx]
+        ).all()
 
     def test_indexed_properties(self, td_name, device):
         td = getattr(self, td_name)(device)
@@ -4459,7 +4463,7 @@ def test_lock_write(self, td_name, device):
         td.lock_()
         td_clone = td.clone()
         assert not td_clone.is_locked
-        td_clone = td.to_tensordict()
+        td_clone = td.to_tensordict(retain_none=True)
         assert not td_clone.is_locked
         assert td.is_locked
         if td_name == "td_h5":
@@ -4752,7 +4756,7 @@ def test_memmap_threads(self, td_name, device, use_dir, tmpdir, num_threads):
     @pytest.mark.parametrize("inplace", [False, True])
     def test_named_apply(self, td_name, device, inplace):
         td = getattr(self, td_name)(device)
-        td_c = td.to_tensordict()
+        td_c = td.to_tensordict(retain_none=True)
 
         def named_plus(name, x):
             if "a" in name:
@@ -4916,7 +4920,7 @@ def test_nestedtensor_stack(self, td_name, device, dim, key):
         with pytest.raises(
             RuntimeError, match="Found more than one unique shape in the tensors"
         ):
-            td_stack.to_tensordict()
+            td_stack.to_tensordict(retain_none=True)
         # cloning is type-preserving: we can do that operation
         td_stack.clone()
 
@@ -5684,7 +5688,7 @@ def test_setitem(self, td_name, device, idx):
             pytest.mark.skip("cannot index tensor with desired index")
             return
 
-        td_clone = td[idx].to_tensordict().zero_()
+        td_clone = td[idx].to_tensordict(retain_none=True).zero_()
         if td_name == "td_params":
             td.data[idx] = td_clone
         else:
@@ -5762,7 +5766,7 @@ def test_setitem_slice(self, td_name, device):
             td_set = td.data
         else:
             td_set = td
-        td_set[:1] = td[:1].to_tensordict().zero_()
+        td_set[:1] = td[:1].to_tensordict(retain_none=True).zero_()
         assert (td[:1] == 0).all()
 
         # with broadcast
@@ -5778,7 +5782,7 @@ def test_setitem_slice(self, td_name, device):
             td_set = td.data
         else:
             td_set = td
-        td_set[:1] = td[0].to_tensordict().zero_()
+        td_set[:1] = td[0].to_tensordict(retain_none=True).zero_()
         assert (td[:1] == 0).all()
 
         td = getattr(self, td_name)(device)
@@ -5793,7 +5797,7 @@ def test_setitem_slice(self, td_name, device):
             td_set = td.data
         else:
             td_set = td
-        td_set[:1, 0] = td[0, 0].to_tensordict().zero_()
+        td_set[:1, 0] = td[0, 0].to_tensordict(retain_none=True).zero_()
         assert (td[:1, 0] == 0).all()
 
         td = getattr(self, td_name)(device)
@@ -5808,7 +5812,7 @@ def test_setitem_slice(self, td_name, device):
             td_set = td.data
         else:
             td_set = td
-        td_set[:1, :, 0] = td[0, :, 0].to_tensordict().zero_()
+        td_set[:1, :, 0] = td[0, :, 0].to_tensordict(retain_none=True).zero_()
         assert (td[:1, :, 0] == 0).all()
 
     def test_setitem_string(self, td_name, device):
@@ -6294,7 +6298,7 @@ def is_namedtuple(obj):
     def test_to_tensordict(self, td_name, device):
         torch.manual_seed(1)
         td = getattr(self, td_name)(device)
-        td2 = td.to_tensordict()
+        td2 = td.to_tensordict(retain_none=True)
         assert (td2 == td).all()
 
     @set_lazy_legacy(True)
@@ -6664,10 +6668,10 @@ def test_update_subtensordict(self, td_name, device, index):
             )
 
         sub_td = td._get_sub_tensordict(index)
-        assert sub_td.shape == td.to_tensordict()[index].shape
+        assert sub_td.shape == td.to_tensordict(retain_none=True)[index].shape
         assert sub_td.shape == td[index].shape, (td, index)
         td0 = td[index]
-        td0 = td0.to_tensordict()
+        td0 = td0.to_tensordict(retain_none=True)
         td0 = td0.apply(lambda x: x * 0 + 2)
         assert sub_td.shape == td0.shape
         if td_name == "td_params":
@@ -6801,7 +6805,7 @@ def test_where_pad(self, td_name, device):
         # test with other empty td
         mask = torch.zeros(td.shape, dtype=torch.bool, device=td.device).bernoulli_()
         if td_name in ("td_h5",):
-            td_full = td.to_tensordict()
+            td_full = td.to_tensordict(retain_none=True)
         else:
             td_full = td
         td_empty = td_full.empty()
@@ -7381,7 +7385,7 @@ def td(self, device):
 
     def td_reset_bs(self, device):
         td = self.td(device)
-        td = td.unsqueeze(-1).to_tensordict()
+        td = td.unsqueeze(-1).to_tensordict(retain_none=True)
         td.batch_size = torch.Size([3, 1])
         return td
 
@@ -8129,7 +8133,7 @@ def test_lazy_indexing(self, pos1, pos2, pos3):
         inner = LazyStackedTensorDict.lazy_stack([td_leaf_1] * 4, 0)
         middle = LazyStackedTensorDict.lazy_stack([inner] * 3, 0)
         outer = LazyStackedTensorDict.lazy_stack([middle] * 2, 0)
-        outer_dense = outer.to_tensordict()
+        outer_dense = outer.to_tensordict(retain_none=True)
         pos1 = self._idx_list[pos1]
         pos2 = self._idx_list[pos2]
         pos3 = self._idx_list[pos3]
@@ -8582,13 +8586,13 @@ def test_stack_update_heter_stacked_td(self, stack_dim):
             RuntimeError,
             match="Found more than one unique shape in the tensors to be stacked",
         ):
-            td_a.update(td_b.to_tensordict())
+            td_a.update(td_b.to_tensordict(retain_none=True))
         td_a.update_(td_b)
         with pytest.raises(
             RuntimeError,
             match="Found more than one unique shape in the tensors to be stacked",
         ):
-            td_a.update_(td_b.to_tensordict())
+            td_a.update_(td_b.to_tensordict(retain_none=True))
 
     def test_stack_with_heterogeneous_stacks(self):
         # tests that we can stack several stacks in a dense manner
@@ -8867,7 +8871,7 @@ def test_inplace(self, save_name):
         path = f"/tmp/{uuid.uuid4()}"
         snapshot = torchsnapshot.Snapshot.take(app_state=app_state, path=path)
 
-        td_plain = td.to_tensordict()
+        td_plain = td.to_tensordict(retain_none=True)
         # we want to delete refs to MemoryMappedTensors
         assert not isinstance(td_plain["a"], MemoryMappedTensor)
         del td
@@ -8896,7 +8900,7 @@ def test_update(self):
         tensordict.memmap_()
         path = f"/tmp/{uuid.uuid4()}"
         snapshot = torchsnapshot.Snapshot.take(app_state=state, path=path)
-        td_plain = tensordict.to_tensordict()
+        td_plain = tensordict.to_tensordict(retain_none=True)
         assert not isinstance(td_plain["a"], MemoryMappedTensor)
         del tensordict
 
@@ -10149,8 +10153,12 @@ def test_index_with_generator(self, chunksize, num_chunks, h5, has_out, tmpdir):
             assert input.shape == input_h5.shape
             input = input_h5
         if has_out:
-            output_generator = torch.zeros_like(self.selectfn(input.to_tensordict()))
-            output_split = torch.zeros_like(self.selectfn(input.to_tensordict()))
+            output_generator = torch.zeros_like(
+                self.selectfn(input.to_tensordict(retain_none=True))
+            )
+            output_split = torch.zeros_like(
+                self.selectfn(input.to_tensordict(retain_none=True))
+            )
         else:
             output_generator = None
             output_split = None