[Feature] Truly invertible tensordict permutation of dimensions (#295)

Author: ramonmedel

test/test_tensordict.py

@@ -306,6 +306,24 @@ def test_permute(device):
     assert torch.sum(t["a"]) == torch.Tensor([0])
 
 
+@pytest.mark.parametrize("device", get_available_devices())
+def test_permute_applied_twice(device):
+    torch.manual_seed(1)
+    d = {
+        "a": torch.randn(4, 5, 6, 9, device=device),
+        "b": torch.randn(4, 5, 6, 7, device=device),
+        "c": torch.randn(4, 5, 6, device=device),
+    }
+    td1 = TensorDict(batch_size=(4, 5, 6), source=d)
+    td2 = torch.permute(td1, dims=(2, 1, 0))
+    td3 = torch.permute(td2, dims=(2, 1, 0))
+    assert td3 is td1
+    td1 = TensorDict(batch_size=(4, 5, 6), source=d)
+    td2 = torch.permute(td1, dims=(2, 1, 0))
+    td3 = torch.permute(td2, dims=(0, 1, 2))
+    assert td3 is not td1
+
+
 @pytest.mark.parametrize("device", get_available_devices())
 def test_permute_exceptions(device):
     torch.manual_seed(1)
@@ -626,6 +644,18 @@ def permute_td(self, device):
         #     batch_size=[3, 1, 2, 4],
         # ).permute(2, 0, 1, 3)
 
+    def test_permute_applied_twice(self, td_name, device):
+        torch.manual_seed(0)
+        tensordict = getattr(self, td_name)(device)
+        for _ in range(10):
+            p = torch.randperm(4)
+            inv_p = p.argsort()
+            other_p = inv_p
+            while (other_p == inv_p).all():
+                other_p = torch.randperm(4)
+            assert tensordict.permute(*p).permute(*inv_p) is tensordict
+            assert tensordict.permute(*p).permute(*other_p) is not tensordict
+
     def unsqueezed_td(self, device):
         td = TensorDict(
             source={

torchrl/data/tensordict/tensordict.py

@@ -1291,6 +1291,10 @@ def permute(
                 f"number of dims don't match in permute (got {len(dims_list)}, expected {len(self.shape)}"
             )
 
+        if not len(dims_list) and not self.batch_dims:
+            return self
+        if np.array_equal(dims_list, range(self.batch_dims)):
+            return self
         min_dim, max_dim = -self.batch_dims, self.batch_dims - 1
         seen = [False for dim in range(max_dim + 1)]
         for idx in dims_list:
@@ -4271,8 +4275,45 @@ def view(
 class PermutedTensorDict(_CustomOpTensorDict):
     """
     A lazy view on a TensorDict with the batch dimensions permuted.
+
+    When calling `tensordict.permute(dims_list, dim)`, a lazy view of this operation is
+    returned such that the following code snippet works without raising an
+    exception:
+
+        >>> assert tensordict.permute(dims_list, dim).permute(dims_list, dim) is tensordict
+
+    Examples:
+        >>> from torchrl.data import TensorDict
+        >>> import torch
+        >>> td = TensorDict({'a': torch.randn(4, 5, 6, 9)}, batch_size=[3])
+        >>> td_permute = td.permute(dims=(2, 1, 0))
+        >>> print(td_permute.shape)
+        torch.Size([6, 5, 4])
+        >>> print(td_permute.permute(dims=(2, 1, 0)) is td)
+        True
     """
 
+    def permute(
+        self,
+        *dims_list: int,
+        dims=None,
+    ) -> _TensorDict:
+        if len(dims_list) == 0:
+            dims_list = dims
+        elif len(dims_list) == 1 and not isinstance(dims_list[0], int):
+            dims_list = dims_list[0]
+        if len(dims_list) != len(self.shape):
+            raise RuntimeError(
+                f"number of dims don't match in permute (got {len(dims_list)}, expected {len(self.shape)}"
+            )
+        if not len(dims_list) and not self.batch_dims:
+            return self
+        if np.array_equal(dims_list, range(self.batch_dims)):
+            return self
+        if np.array_equal(np.argsort(dims_list), self.inv_op_kwargs.get("dims")):
+            return self._source
+        return super().permute(*dims_list)
+
     def add_missing_dims(self, num_dims: int, batch_dims: tuple) -> tuple:
         dim_diff = num_dims - len(batch_dims)
         all_dims = [i for i in range(num_dims)]