[Feature] flatten_keys and unflatten_keys as context managers

tensordict/_td.py

@@ -2952,7 +2952,7 @@ def keys(
         leaves_only: bool = False,
         is_leaf: Callable[[Type], bool] | None = None,
     ) -> _TensorDictKeysView:
-        if not include_nested and not leaves_only:
+        if not include_nested and not leaves_only and is_leaf is None:
             return _StringKeys(self._tensordict.keys())
         else:
             return self._nested_keys(

tensordict/base.py

@@ -7931,6 +7931,18 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                     return out.update(self.transpose(dim0, dim1), inplace=False)
                 else:
                     return out.update_(self.transpose(dim0, dim1))
+            elif last_op == type(self).flatten_keys.__name__:
+                sep = args[0] if args else "."
+                if not out.is_locked:
+                    return out.update(self.unflatten_keys(sep), inplace=False)
+                else:
+                    return out.update_(self.unflatten_keys(sep))
+            elif last_op == type(self).unflatten_keys.__name__:
+                sep = args[0] if args else "."
+                if not out.is_locked:
+                    return out.update(self.flatten_keys(sep), inplace=False)
+                else:
+                    return out.update_(self.flatten_keys(sep))
             elif last_op == type(self).flatten.__name__:
                 if len(args) == 2:
                     dim0, dim1 = args
@@ -8623,6 +8635,7 @@ def contiguous(self) -> T:
         ...
 
     @cache  # noqa: B019
+    @_as_context_manager()
     def flatten_keys(
         self,
         separator: str = ".",
@@ -8778,6 +8791,7 @@ def _flatten_keys_inplace(self, separator, is_leaf):
         return self
 
     @cache  # noqa: B019
+    @_as_context_manager()
     def unflatten_keys(self, separator: str = ".", inplace: bool = False) -> T:
         """Converts a flat tensordict into a nested one, recursively.
 

tensordict/persistent.py

@@ -36,6 +36,7 @@
 )
 from tensordict.memmap import MemoryMappedTensor
 from tensordict.utils import (
+    _as_context_manager,
     _CloudpickleWrapper,
     _KEY_ERROR,
     _LOCK_ERROR,
@@ -943,13 +944,15 @@ def _exclude(
             "Create a regular tensordict first using the `to_tensordict` method."
         )
 
+    @_as_context_manager()
     def flatten_keys(self, separator: str = ".", inplace: bool = False) -> T:
         if inplace:
             raise ValueError(
                 "Cannot call flatten_keys in_place with a PersistentTensorDict."
             )
         return self.to_tensordict().flatten_keys(separator=separator)
 
+    @_as_context_manager()
     def unflatten_keys(self, separator: str = ".", inplace: bool = False) -> T:
         if inplace:
             raise ValueError(

test/test_tensordict.py

@@ -3514,6 +3514,21 @@ def test_flatten_keys(self, td_name, device, inplace, separator):
         else:
             assert td_flatten is not td
 
+    def test_flatten_keys_decorator(self, td_name, device):
+        td = getattr(self, td_name)(device)
+        with td.flatten_keys(",") as tdflat:
+            assert set(tdflat.keys(True, True, is_leaf=_is_leaf_nontensor)) == set(
+                tdflat.keys(is_leaf=_is_leaf_nontensor)
+            )
+            with tdflat.unflatten_keys(",") as td_orig:
+                assert (td_orig == td).all()
+                if not td.is_locked:
+                    td_orig["new", "data"] = torch.zeros(td_orig.shape)
+            if not td.is_locked:
+                assert (tdflat["new,data"] == 0).all()
+        if not td.is_locked:
+            assert (td["new", "data"] == 0).all()
+
     def test_flatten_unflatten(self, td_name, device):
         td = getattr(self, td_name)(device)
         shape = td.shape[:3]