[Feature] flatten and unflatten as decorators

tensordict/_td.py

@@ -1866,7 +1866,16 @@ def _set_str(
                     dest.update(value, inplace=True, non_blocking=non_blocking)
                 else:
                     if dest is not value:
-                        dest.copy_(value, non_blocking=non_blocking)
+                        try:
+                            dest.copy_(value, non_blocking=non_blocking)
+                        except RuntimeError:
+                            # if we're updating a param and the storages match, nothing needs to be done
+                            if not (
+                                isinstance(dest, torch.Tensor)
+                                and dest.data.untyped_storage().data_ptr()
+                                == value.data.untyped_storage().data_ptr()
+                            ):
+                                raise
             except KeyError as err:
                 raise err
             except Exception as err:

tensordict/base.py

@@ -3840,6 +3840,7 @@ def sorted_keys(self) -> list[NestedKey]:
         """
         return sorted(self.keys())
 
+    @as_decorator()
     def flatten(self, start_dim=0, end_dim=-1):
         """Flattens all the tensors of a tensordict.
 
@@ -3871,6 +3872,8 @@ def flatten(self, start_dim=0, end_dim=-1):
             tensor([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11])
 
         """
+        if start_dim < 0:
+            start_dim = self.ndim + start_dim
         if end_dim < 0:
             end_dim = self.ndim + end_dim
             if end_dim < 0:
@@ -3906,6 +3909,7 @@ def flatten(tensor):
             out.names = names
         return out
 
+    @as_decorator()
     def unflatten(self, dim, unflattened_size):
         """Unflattens a tensordict dim expanding it to a desired shape.
 
@@ -6128,21 +6132,71 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         _last_op = self._last_op_queue.pop()
         if _last_op is not None:
             last_op, (args, kwargs, out) = _last_op
+            # TODO: transpose, flatten etc. as decorator should lock the content to make sure that no key is
+            #  added or deleted
             if last_op == self.__class__.lock_.__name__:
                 return self.unlock_()
             elif last_op == self.__class__.unlock_.__name__:
                 return self.lock_()
             elif last_op == self.__class__.transpose.__name__:
                 dim0, dim1 = args
-                return out.update(self.transpose(dim0, dim1))
+                if not out.is_locked:
+                    return out.update(self.transpose(dim0, dim1), inplace=True)
+                else:
+                    return out.update_(self.transpose(dim0, dim1))
+            elif last_op == self.__class__.flatten.__name__:
+                if len(args) == 2:
+                    dim0, dim1 = args
+                elif len(args) == 1:
+                    dim0 = args[0]
+                    dim1 = kwargs.get("end_dim", -1)
+                else:
+                    dim0 = kwargs.get("start_dim", 0)
+                    dim1 = kwargs.get("end_dim", -1)
+                if dim1 < 0:
+                    dim1 = out.ndim + dim1
+                if dim0 < 0:
+                    dim0 = out.ndim + dim0
+
+                if not out.is_locked:
+                    return out.update(
+                        self.unflatten(dim0, out.shape[dim0 : dim1 + 1]), inplace=True
+                    )
+                else:
+                    return out.update_(self.unflatten(dim0, out.shape[dim0 : dim1 + 1]))
+
+            elif last_op == self.__class__.unflatten.__name__:
+                if args:
+                    dim0 = args[0]
+                    if len(args) > 1:
+                        unflattened_size = args[1]
+                    else:
+                        unflattened_size = kwargs.get("unflattened_size")
+                else:
+                    dim0 = kwargs.get("dim")
+                    unflattened_size = kwargs.get("unflattened_size")
+                if dim0 < 0:
+                    dim0 = out.ndim + dim0
+                dim1 = dim0 + len(unflattened_size) - 1
+                if not out.is_locked:
+                    return out.update(self.flatten(dim0, dim1), inplace=True)
+                else:
+                    return out.update_(self.flatten(dim0, dim1))
+
             elif last_op == self.__class__.permute.__name__:
                 dims_list = _get_shape_from_args(*args, kwarg_name="dims", **kwargs)
                 dims_list = [dim if dim >= 0 else self.ndim + dim for dim in dims_list]
                 # inverse map
                 inv_dims_list = np.argsort(dims_list)
-                return out.update(self.permute(inv_dims_list))
+                if not out.is_locked:
+                    return out.update(self.permute(inv_dims_list), inplace=True)
+                else:
+                    return out.update_(self.permute(inv_dims_list))
             elif last_op == self.__class__.view.__name__:
-                return out.update(self.view(out.shape))
+                if not out.is_locked:
+                    return out.update(self.view(out.shape), inplace=True)
+                else:
+                    return out.update_(self.view(out.shape))
             elif last_op == self.__class__.unsqueeze.__name__:
                 if args:
                     (dim,) = args
@@ -6152,7 +6206,10 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                     raise RuntimeError(
                         "Cannot use td.unsqueeze() as a decorator if the dimension is implicit."
                     )
-                return out.update(self.squeeze(dim))
+                if not out.is_locked:
+                    return out.update(self.squeeze(dim), inplace=True)
+                else:
+                    return out.update_(self.squeeze(dim))
             elif last_op == self.__class__.squeeze.__name__:
                 if args:
                     (dim,) = args
@@ -6162,7 +6219,10 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                     raise RuntimeError(
                         "Cannot use td.squeeze() as a decorator if the dimension is implicit."
                     )
-                return out.update(self.unsqueeze(dim))
+                if not out.is_locked:
+                    return out.update(self.unsqueeze(dim), inplace=True)
+                else:
+                    return out.update_(self.unsqueeze(dim))
             elif last_op == self.__class__.to_module.__name__:
                 if is_tensor_collection(out):
                     with out.unlock_():

test/test_tensordict.py

@@ -3191,6 +3191,35 @@ def test_flatten_unflatten(self, td_name, device):
         assert (td.to_tensordict() == td_unflat).all()
         assert td.batch_size == td_unflat.batch_size
 
+    @pytest.mark.parametrize("start_dim", [0, 1, -2, -3])
+    def test_flatten_unflatten_decorator(self, td_name, device, start_dim):
+        td = getattr(self, td_name)(device)
+        with td.unlock_(), td.flatten(start_dim=start_dim, end_dim=3) as td_flat:
+            assert (td_flat == td.flatten(start_dim, 3)).all()
+            new_start_dim = -1 if start_dim in (-2, -3) else start_dim
+            with td_flat.unflatten(
+                dim=new_start_dim, unflattened_size=td.shape[start_dim:]
+            ) as td_unflat:
+                assert (td_unflat == td).all()
+
+        with td.unlock_(), td.flatten(start_dim, end_dim=3) as td_flat:
+            assert (td_flat == td.flatten(start_dim, 3)).all()
+            new_start_dim = (
+                -1 if start_dim == -2 else -1 if start_dim == -3 else start_dim
+            )
+            with td_flat.unflatten(
+                new_start_dim, unflattened_size=td.shape[start_dim:]
+            ) as td_unflat:
+                assert (td_unflat == td).all()
+
+        with td.unlock_(), td.flatten(start_dim, -1) as td_flat:
+            assert (td_flat == td.flatten(start_dim, -1)).all()
+            new_start_dim = (
+                -1 if start_dim == -2 else -1 if start_dim == -3 else start_dim
+            )
+            with td_flat.unflatten(new_start_dim, td.shape[start_dim:]) as td_unflat:
+                assert (td_unflat == td).all()
+
     def test_flatten_unflatten_bis(self, td_name, device):
         td = getattr(self, td_name)(device)
         shape = td.shape[1:4]