[Feature] Adding additional checks to TensorDict.view to remove unnecessary ViewedTensorDict object creation (#319)

Author: bamaxw

test/test_tensordict.py

@@ -488,6 +488,16 @@ def test_savedtensordict(device):
     assert ss.get("a").device == device
 
 
+def test_inferred_view_size():
+    td = TensorDict({"a": torch.randn(3, 4)}, [3, 4])
+    assert td.view(-1).view(-1, 4) is td
+
+    assert td.view(-1, 4) is td
+    assert td.view(3, -1) is td
+    assert td.view(3, 4) is td
+    assert td.view(-1, 12).shape == torch.Size([1, 12])
+
+
 @pytest.mark.parametrize(
     "ellipsis_index, expected_index",
     [
@@ -1026,6 +1036,22 @@ def test_view(self, td_name, device):
         assert (td_view.get("a") == 1).all()
         assert (td.get("a") == 1).all()
 
+    def test_inferred_view_size(self, td_name, device):
+        if td_name in ("permute_td", "sub_td2"):
+            pytest.skip("view incompatible with stride / permutation")
+        torch.manual_seed(1)
+        td = getattr(self, td_name)(device)
+        for i in range(len(td.shape)):
+            # replacing every index one at a time
+            # with -1, to test that td.view(..., -1, ...)
+            # always returns the original tensordict
+            new_shape = [
+                dim_size if dim_idx != i else -1
+                for dim_idx, dim_size in enumerate(td.shape)
+            ]
+            assert td.view(-1).view(*new_shape) is td
+            assert td.view(*new_shape) is td
+
     def test_clone_td(self, td_name, device):
         torch.manual_seed(1)
         td = getattr(self, td_name)(device)

torchrl/data/tensordict/memmap.py

@@ -52,7 +52,7 @@ class MemmapTensor(object):
     and remote workers that have access to
     a common storage, and as such it supports serialization and
     deserialization. It is possible to choose if the ownership is
-    transferred upon serialization / deserialization: If owenership is not
+    transferred upon serialization / deserialization: If ownership is not
     transferred (transfer_ownership=False, default), then the process where
     the MemmapTensor was created will be responsible of clearing it once it
     gets out of scope (in that process). Otherwise, the process that

torchrl/data/tensordict/tensordict.py

@@ -33,6 +33,7 @@
 
 import numpy as np
 import torch
+from torch.jit._shape_functions import infer_size_impl
 
 from torchrl import KeyDependentDefaultDict, prod
 from torchrl.data.tensordict.memmap import MemmapTensor
@@ -1232,7 +1233,10 @@ def view(
         elif len(shape) == 1 and isinstance(shape[0], (list, tuple, torch.Size)):
             return self.view(*shape[0])
         elif not isinstance(shape, torch.Size):
+            shape = infer_size_impl(shape, self.numel())
             shape = torch.Size(shape)
+        if shape == self.shape:
+            return self
         return ViewedTensorDict(
             source=self,
             custom_op="view",
@@ -4298,6 +4302,7 @@ def view(
         elif len(shape) == 1 and isinstance(shape[0], (list, tuple, torch.Size)):
             return self.view(*shape[0])
         elif not isinstance(shape, torch.Size):
+            shape = infer_size_impl(shape, self.numel())
             shape = torch.Size(shape)
         if shape == self._source.batch_size:
             return self._source