Repository: pytorch/tensordict

Issue:
[BUG] Conflicting keys silently collide in `tensordict.unflatten_keys`
## Describe the bug

If two keys have the same resulting index in `flatten_keys` and `unflatten_keys`, they will be overwritten without an error being raised. This may lead to unexpected behaviours down the line.

## To Reproduce

We would expect this code to fail:
```python
from tensordict import TensorDict
import torch
t = TensorDict({'a.a': torch.randn(3), 'a': {'a': torch.randn(3)}}, [])
print(t.unflatten_keys('.'))
```
as `("a", "a")` already exists.
Similarly, 
```
from tensordict import TensorDict
import torch
t = TensorDict({'a.a': torch.randn(3), 'a': {'a': torch.randn(3)}}, [])
print(t.flatten_keys('.'))
```
should also fail as `"a.a"` is also an existing key.


Discussion:
working on this

Resolving patch:
diff --git a/tensordict/tensordict.py b/tensordict/tensordict.py
--- a/tensordict/tensordict.py
+++ b/tensordict/tensordict.py
@@ -1619,9 +1619,19 @@ def flatten_keys(
         self, separator: str = ".", inplace: bool = False
     ) -> TensorDictBase:
         to_flatten = []
+        existing_keys = self.keys(include_nested=True)
         for key, meta_value in self.items_meta():
+            key_split = tuple(key.split(separator))
             if meta_value.is_tensordict():
                 to_flatten.append(key)
+            elif (
+                separator in key
+                and key_split in existing_keys
+                and not self._get_meta(key_split).is_tensordict()
+            ):
+                raise KeyError(
+                    f"Flattening keys in tensordict collides with existing key '{key}'"
+                )
 
         if inplace:
             for key in to_flatten:
@@ -1670,7 +1680,14 @@ def unflatten_keys(
         else:
             out = self
 
+        keys = set(out.keys())
         for key, list_of_keys in to_unflatten.items():
+
+            if key in keys:
+                raise KeyError(
+                    "Unflattening key(s) in tensordict will override existing unflattened key"
+                )
+
             tensordict = TensorDict({}, batch_size=self.batch_size, device=self.device)
             if key in self.keys():
                 tensordict.update(self[key])


Test patch:
diff --git a/test/test_tensordict.py b/test/test_tensordict.py
--- a/test/test_tensordict.py
+++ b/test/test_tensordict.py
@@ -2926,13 +2926,106 @@ def test_update_nested_dict():
     t.update({"a": {"d": [[[1]] * 3] * 2}})
 
 
-@pytest.mark.parametrize("separator", [".", "-"])
-def test_unflatten_keys_collision(separator):
-    td = TensorDict(
-        {"a": [1, 2], f"c{separator}a": [1, 2], "c": TensorDict({"b": [1, 2]}, [])}, []
+@pytest.mark.parametrize("inplace", [True, False])
+@pytest.mark.parametrize("separator", [",", "-"])
+def test_flatten_unflatten_key_collision(inplace, separator):
+    td1 = TensorDict(
+        {
+            f"a{separator}b{separator}c": torch.zeros(3),
+            "a": {"b": {"c": torch.zeros(3)}},
+        },
+        [],
     )
-    ref = TensorDict({"a": [1, 2], "c": TensorDict({"a": [1, 2], "b": [1, 2]}, [])}, [])
-    assert assert_allclose_td(td.unflatten_keys(separator), ref)
+    td2 = TensorDict(
+        {
+            f"a{separator}b": torch.zeros(3),
+            "a": {"b": torch.zeros(3)},
+            "g": {"d": torch.zeros(3)},
+        },
+        [],
+    )
+    td3 = TensorDict(
+        {
+            f"a{separator}b{separator}c": torch.zeros(3),
+            "a": {"b": {"c": torch.zeros(3), "d": torch.zeros(3)}},
+        },
+        [],
+    )
+
+    td4 = TensorDict(
+        {
+            f"a{separator}b{separator}c{separator}d": torch.zeros(3),
+            "a": {"b": {"c": torch.zeros(3)}},
+        },
+        [],
+    )
+
+    td5 = TensorDict(
+        {f"a{separator}b": torch.zeros(3), "a": {"b": {"c": torch.zeros(3)}}}, []
+    )
+
+    with pytest.raises(
+        KeyError, match="Flattening keys in tensordict collides with existing key *"
+    ):
+        _ = td1.flatten_keys(separator)
+
+    with pytest.raises(
+        KeyError, match="Flattening keys in tensordict collides with existing key *"
+    ):
+        _ = td2.flatten_keys(separator)
+
+    with pytest.raises(
+        KeyError, match="Flattening keys in tensordict collides with existing key *"
+    ):
+        _ = td3.flatten_keys(separator)
+
+    with pytest.raises(
+        KeyError,
+        match=re.escape(
+            "Unflattening key(s) in tensordict will override existing unflattened key"
+        ),
+    ):
+        _ = td1.unflatten_keys(separator)
+
+    with pytest.raises(
+        KeyError,
+        match=re.escape(
+            "Unflattening key(s) in tensordict will override existing unflattened key"
+        ),
+    ):
+        _ = td2.unflatten_keys(separator)
+
+    with pytest.raises(
+        KeyError,
+        match=re.escape(
+            "Unflattening key(s) in tensordict will override existing unflattened key"
+        ),
+    ):
+        _ = td3.unflatten_keys(separator)
+
+    with pytest.raises(
+        KeyError,
+        match=re.escape(
+            "Unflattening key(s) in tensordict will override existing unflattened key"
+        ),
+    ):
+        _ = td4.unflatten_keys(separator)
+
+    with pytest.raises(
+        KeyError,
+        match=re.escape(
+            "Unflattening key(s) in tensordict will override existing unflattened key"
+        ),
+    ):
+        _ = td5.unflatten_keys(separator)
+
+    td4_flat = td4.flatten_keys(separator)
+    assert (f"a{separator}b{separator}c{separator}d") in td4_flat.keys()
+    assert (f"a{separator}b{separator}c") in td4_flat.keys()
+
+    td5_flat = td5.flatten_keys(separator)
+    assert (f"a{separator}b") in td5_flat.keys()
+    assert (f"a{separator}b{separator}c") in td5_flat.keys()
 
 
 if __name__ == "__main__":