Repository: pytorch/tensordict

Issue:
[BUG] __repr__ in nested lazy stacks
Tor reproduce
```
def get_agent_tensors(
    i,
):
    camera = torch.zeros(32, 32, 3)
    vector_3d = torch.zeros(3)
    vector_2d = torch.zeros(2)
    lidar = torch.zeros(20)

    agent_0_obs = torch.zeros(1)
    agent_1_obs = torch.zeros(1, 2)
    agent_2_obs = torch.zeros(1, 2, 3)

    # Agents all have the same camera
    # All have vector entry but different shapes
    # First 2 have lidar and last sonar
    # All have a different key agent_i_obs with different n_dims
    if i == 0:
        return TensorDict(
            {
                "camera": camera,
                "lidar": lidar,
                "vector": vector_3d,
                "agent_0_obs": agent_0_obs,
            },
            [],
        )
    elif i == 1:
        return TensorDict(
            {
                "camera": camera,
                "lidar": lidar,
                "vector": vector_2d,
                "agent_1_obs": agent_1_obs,
            },
            [],
        )
    elif i == 2:
        return TensorDict(
            {
                "camera": camera,
                "vector": vector_2d,
                "agent_2_obs": agent_2_obs,
            },
            [],
        )
    else:
        raise ValueError(f"Index {i} undefined for 3 agents")


@pytest.mark.parametrize("batch_size", [(), (32,), (32, 4)])
def test_lazy_stack_stack(batch_size):
    agent_obs = []
    for angent_id in range(3):
        agent_obs.append(get_agent_tensors(angent_id))
    agent_obs = torch.stack(agent_obs, dim=0)
    obs = TensorDict(
        {
            "agents": agent_obs,
            "state": torch.zeros(
                64,
                64,
                3,
            ),
        },
        [],
    )
    obs = obs.expand(batch_size)


    obs_stack = torch.stack([obs])
    print(obs_stack["agents"])


    obs2 = obs.clone()
    obs_stack = torch.stack([obs, obs2])
    print(obs_stack["agents"])

```

I can try to work on a fix and integrate in #493


Discussion:

Resolving patch:
diff --git a/tensordict/tensordict.py b/tensordict/tensordict.py
--- a/tensordict/tensordict.py
+++ b/tensordict/tensordict.py
@@ -3243,7 +3243,6 @@ def is_boolean(idx):
             if num_boolean_dim:
                 names = [None] + names[num_boolean_dim:]
             else:
-
                 # def is_int(subidx):
                 #     if isinstance(subidx, Number):
                 #         return True
@@ -3344,7 +3343,6 @@ def __setitem__(
         index: IndexType,
         value: TensorDictBase | dict | numbers.Number | CompatibleType,
     ) -> None:
-
         istuple = isinstance(index, tuple)
         if istuple or isinstance(index, str):
             # try:
@@ -4733,7 +4731,6 @@ def _gather_tensor(tensor, dest=None):
         return out
 
     if out is None:
-
         names = input.names if input._has_names() else None
 
         return TensorDict(
@@ -6439,7 +6436,6 @@ def _get_str(
         key: NestedKey,
         default: str | CompatibleType = NO_DEFAULT,
     ) -> CompatibleType:
-
         # we can handle the case where the key is a tuple of length 1
         tensors = []
         for td in self.tensordicts:
@@ -7107,7 +7103,6 @@ def _isend(
         _futures: list[torch.Future] | None = None,
         pseudo_rand: bool = False,
     ) -> int:
-
         if _futures is None:
             is_root = True
             _futures = []
@@ -7184,7 +7179,6 @@ def del_(self, key: str, **kwargs: Any) -> TensorDictBase:
     def pop(
         self, key: NestedKey, default: str | CompatibleType = NO_DEFAULT
     ) -> CompatibleType:
-
         # using try/except for get/del is suboptimal, but
         # this is faster that checkink if key in self keys
         key = _unravel_key_to_tuple(key)
@@ -7566,6 +7560,44 @@ def __del__(self):
         for td in self.tensordicts:
             td._remove_lock(id(self))
 
+    def __repr__(self):
+        fields = _td_fields(self)
+        field_str = indent(f"fields={{{fields}}}", 4 * " ")
+        exclusive_fields_str = indent(
+            f"exclusive_fields={{{self._repr_exclusive_fields()}}}", 4 * " "
+        )
+        batch_size_str = indent(f"batch_size={self.batch_size}", 4 * " ")
+        device_str = indent(f"device={self.device}", 4 * " ")
+        is_shared_str = indent(f"is_shared={self.is_shared()}", 4 * " ")
+        stack_dim = indent(f"stack_dim={self.stack_dim}", 4 * " ")
+        string = ",\n".join(
+            [
+                field_str,
+                exclusive_fields_str,
+                batch_size_str,
+                device_str,
+                is_shared_str,
+                stack_dim,
+            ]
+        )
+        return f"{type(self).__name__}(\n{string})"
+
+    def _repr_exclusive_fields(self):
+        keys = set(self.keys())
+        exclusive_keys = [
+            _td_fields(td, [k for k in td.keys() if k not in keys])
+            for td in self.tensordicts
+        ]
+        exclusive_key_str = ",\n".join(
+            [
+                indent(f"{i} ->{line}", 4 * " ")
+                for i, line in enumerate(exclusive_keys)
+                if line != "\n"
+            ]
+        )
+
+        return "\n" + exclusive_key_str
+
     lock_ = TensorDictBase.lock_
     lock = _renamed_inplace_method(lock_)
 
@@ -8169,7 +8201,6 @@ def _stack_onto_(
         list_item: list[CompatibleType],
         dim: int,
     ) -> TensorDictBase:
-
         trsp = self.custom_op_kwargs["dim0"], self.custom_op_kwargs["dim1"]
         if dim == trsp[0]:
             dim = trsp[1]
@@ -8336,31 +8367,35 @@ def _make_repr(key: str, item: CompatibleType, tensordict: TensorDictBase) -> st
     return f"{key}: {_get_repr(item)}"
 
 
-def _td_fields(td: TensorDictBase) -> str:
+def _td_fields(td: TensorDictBase, keys=None) -> str:
     strs = []
-    for key in td.keys():
-        try:
+    if keys is None:
+        keys = td.keys()
+    for key in keys:
+        shape = td.get_item_shape(key)
+        if -1 not in shape:
             item = td.get(key)
             strs.append(_make_repr(key, item, td))
-        except RuntimeError as err:
-            if re.match(r"Found more than one unique shape in the tensors", str(err)):
-                # we know td is lazy stacked and the key is a leaf
-                # so we can get the shape and escape the error
-                shape = td.get_item_shape(key)
-                tensor = td.tensordicts[0].get(key)
-                if isinstance(tensor, TensorDictBase):
-                    substr = _td_fields(tensor)
-                else:
-                    substr = _get_repr_custom(
-                        tensor.__class__,
-                        shape=shape,
-                        device=tensor.device,
-                        dtype=tensor.dtype,
-                        is_shared=tensor.is_shared(),
-                    )
-                strs.append(f"{key}: {substr}")
+        else:
+            # we know td is lazy stacked and the key is a leaf
+            # so we can get the shape and escape the error
+            temp_td = td
+            while isinstance(
+                temp_td, LazyStackedTensorDict
+            ):  # we need to grab the het tensor from the inner nesting level
+                temp_td = temp_td.tensordicts[0]
+            tensor = temp_td.get(key)
+            if isinstance(tensor, TensorDictBase):
+                substr = _td_fields(tensor)
             else:
-                raise err
+                substr = _get_repr_custom(
+                    tensor.__class__,
+                    shape=shape,
+                    device=tensor.device,
+                    dtype=tensor.dtype,
+                    is_shared=tensor.is_shared(),
+                )
+            strs.append(f"{key}: {substr}")
 
     return indent(
         "\n" + ",\n".join(sorted(strs)),


Test patch:
diff --git a/test/test_tensordict.py b/test/test_tensordict.py
--- a/test/test_tensordict.py
+++ b/test/test_tensordict.py
@@ -2721,9 +2721,15 @@ def test_repr_stacked(self, device, dtype):
         expected = f"""LazyStackedTensorDict(
     fields={{
         a: {tensor_class}(shape=torch.Size([4, 3, 2, 1, 5]), device={tensor_device}, dtype={dtype}, is_shared={is_shared_tensor})}},
+    exclusive_fields={{
+        0 ->
+            c: {tensor_class}(shape=torch.Size([4, 3, 1, 5]), device={tensor_device}, dtype={dtype}, is_shared={is_shared_tensor}),
+        1 ->
+            b: {tensor_class}(shape=torch.Size([4, 3, 1, 10]), device={tensor_device}, dtype={dtype}, is_shared={is_shared_tensor})}},
     batch_size=torch.Size([4, 3, 2, 1]),
     device={str(device)},
-    is_shared={is_shared})"""
+    is_shared={is_shared},
+    stack_dim={stacked_td.stack_dim})"""
         assert repr(stacked_td) == expected
 
     def test_repr_stacked_het(self, device, dtype):
@@ -2760,9 +2766,12 @@ def test_repr_stacked_het(self, device, dtype):
     fields={{
         a: Tensor(shape=torch.Size([2, -1]), device={tensor_device}, dtype={dtype}, is_shared={is_shared_tensor}),
         b: Tensor(shape=torch.Size([-1]), device={tensor_device}, dtype={dtype}, is_shared={is_shared_tensor})}},
+    exclusive_fields={{
+    }},
     batch_size=torch.Size([2]),
     device={str(device)},
-    is_shared={is_shared})"""
+    is_shared={is_shared},
+    stack_dim={stacked_td.stack_dim})"""
         assert repr(stacked_td) == expected
 
     @pytest.mark.parametrize("index", [None, (slice(None), 0)])
@@ -2849,20 +2858,20 @@ def test_repr_indexed_stacked_tensordict(self, device, dtype, index):
             is_shared_tensor = True
         else:
             is_shared_tensor = is_shared
-        if index is None:
-            expected = f"""LazyStackedTensorDict(
-    fields={{
-        a: {tensor_class}(shape=torch.Size([4, 3, 2, 1, 5]), device={tensor_device}, dtype={dtype}, is_shared={is_shared_tensor})}},
-    batch_size=torch.Size([4, 3, 2, 1]),
-    device={str(device)},
-    is_shared={is_shared})"""
-        else:
-            expected = f"""LazyStackedTensorDict(
+
+        expected = f"""LazyStackedTensorDict(
     fields={{
         a: {tensor_class}(shape=torch.Size([4, 3, 2, 1, 5]), device={tensor_device}, dtype={dtype}, is_shared={is_shared_tensor})}},
+    exclusive_fields={{
+        0 ->
+            c: {tensor_class}(shape=torch.Size([4, 3, 1, 5]), device={tensor_device}, dtype={dtype}, is_shared={is_shared_tensor}),
+        1 ->
+            b: {tensor_class}(shape=torch.Size([4, 3, 1, 10]), device={tensor_device}, dtype={dtype}, is_shared={is_shared_tensor})}},
     batch_size=torch.Size([4, 3, 2, 1]),
     device={str(device)},
-    is_shared={is_shared})"""
+    is_shared={is_shared},
+    stack_dim={stacked_tensordict.stack_dim})"""
+
         assert repr(stacked_tensordict) == expected
 
     @pytest.mark.skipif(not torch.cuda.device_count(), reason="no cuda")
@@ -4110,6 +4119,85 @@ def test_setitem_hetero(self, batch_size, stack_dim):
         assert self.recursively_check_key(res1[index], 1)  # check all 1
         assert self.recursively_check_key(res2[index], 1)  # check all 1
 
+    @pytest.mark.parametrize("batch_size", [(), (32,), (32, 4)])
+    def test_lazy_stack_stack(self, batch_size):
+        obs = self.nested_lazy_het_td(batch_size)
+
+        assert isinstance(obs, TensorDict)
+        assert isinstance(obs["lazy"], LazyStackedTensorDict)
+        assert obs["lazy"].stack_dim == len(obs["lazy"].shape) - 1  # succeeds
+        assert obs["lazy"].shape == (*batch_size, 3)
+        assert isinstance(obs["lazy"][..., 0], TensorDict)  # succeeds
+
+        obs_stack = torch.stack([obs])
+
+        assert (
+            isinstance(obs_stack, LazyStackedTensorDict) and obs_stack.stack_dim == 0
+        )  # succeeds
+        assert obs_stack.batch_size == (1, *batch_size)  # succeeds
+        assert obs_stack[0] is obs  # succeeds
+        assert isinstance(obs_stack["lazy"], LazyStackedTensorDict)
+        assert obs_stack["lazy"].shape == (1, *batch_size, 3)
+        assert obs_stack["lazy"].stack_dim == 0  # succeeds
+        assert obs_stack["lazy"][0] is obs["lazy"]
+
+        obs2 = obs.clone()
+        obs_stack = torch.stack([obs, obs2])
+
+        assert (
+            isinstance(obs_stack, LazyStackedTensorDict) and obs_stack.stack_dim == 0
+        )  # succeeds
+        assert obs_stack.batch_size == (2, *batch_size)  # succeeds
+        assert obs_stack[0] is obs  # succeeds
+        assert isinstance(obs_stack["lazy"], LazyStackedTensorDict)
+        assert obs_stack["lazy"].shape == (2, *batch_size, 3)
+        assert obs_stack["lazy"].stack_dim == 0  # succeeds
+        assert obs_stack["lazy"][0] is obs["lazy"]
+
+    @pytest.mark.parametrize("batch_size", [(), (32,), (32, 4)])
+    def test_stack_hetero(self, batch_size):
+        obs = self.nested_lazy_het_td(batch_size)
+
+        obs2 = obs.clone()
+        obs2.apply_(lambda x: x + 1)
+
+        obs_stack = torch.stack([obs, obs2])
+        obs_stack_resolved = self.dense_stack_tds_v2([obs, obs2], stack_dim=0)
+
+        assert isinstance(obs_stack, LazyStackedTensorDict) and obs_stack.stack_dim == 0
+        assert isinstance(obs_stack_resolved, TensorDict)
+
+        assert obs_stack.batch_size == (2, *batch_size)
+        assert obs_stack_resolved.batch_size == obs_stack.batch_size
+
+        assert obs_stack["lazy"].shape == (2, *batch_size, 3)
+        assert obs_stack_resolved["lazy"].batch_size == obs_stack["lazy"].batch_size
+
+        assert obs_stack["lazy"].stack_dim == 0
+        assert (
+            obs_stack_resolved["lazy"].stack_dim
+            == len(obs_stack_resolved["lazy"].batch_size) - 1
+        )
+        for stack in [obs_stack_resolved, obs_stack]:
+            for index in range(2):
+                assert (stack[index]["dense"] == index).all()
+                assert (stack["dense"][index] == index).all()
+                assert (stack["lazy"][index]["shared"] == index).all()
+                assert (stack[index]["lazy"]["shared"] == index).all()
+                assert (stack["lazy"]["shared"][index] == index).all()
+                assert (
+                    stack["lazy"][index][..., 0]["individual_0_tensor"] == index
+                ).all()
+                assert (
+                    stack[index]["lazy"][..., 0]["individual_0_tensor"] == index
+                ).all()
+                assert (
+                    stack["lazy"][..., 0]["individual_0_tensor"][index] == index
+                ).all()
+                assert (
+                    stack["lazy"][..., 0][index]["individual_0_tensor"] == index
+                ).all()
+
     def test_add_batch_dim_cache(self):
         td = TensorDict(
             {"a": torch.rand(3, 4, 5), ("b", "c"): torch.rand(3, 4, 5)}, [3, 4, 5]