[BugFix] update TensorDictModule out_keys on assignment

tensordict/nn/common.py

@@ -387,11 +387,21 @@ def __init__(self, out_keys):
     def _init(self, module):
         if self._initialized:
             return
-        self._initialized = True
-        self.module = module
         if not all(key in module.out_keys for key in self.out_keys):
             raise RuntimeError("Some keys are not part of the module out_keys.")
-        module.out_keys = self.out_keys
+        had_out_keys_apparent = "_out_keys_apparent" in module.__dict__
+        out_keys_apparent = module.__dict__.get("_out_keys_apparent")
+        module._out_keys_apparent = self.out_keys
+        if module.out_keys != self.out_keys:
+            if had_out_keys_apparent:
+                module._out_keys_apparent = out_keys_apparent
+            else:
+                del module._out_keys_apparent
+            raise RuntimeError(
+                f"{type(module).__name__} does not support select_out_keys."
+            )
+        self._initialized = True
+        self.module = module
 
     def __call__(  # noqa: F811
         self,
@@ -409,7 +419,7 @@ def __call__(  # noqa: F811
         if not tensordict_in and kwargs.get("tensordict") is not None:
             tensordict_in = kwargs.pop("tensordict")
         is_dispatched = self._detect_dispatch(tensordict_in, kwargs, in_keys)
-        out_keys = self.out_keys
+        out_keys = module.out_keys
         # if dispatch filtered the out keys as they should we're happy
         if is_dispatched:
             if (not isinstance(tensordict_out, tuple) and len(out_keys) == 1) or (
@@ -563,11 +573,7 @@ def out_keys_source(self):
 
     @out_keys.setter
     def out_keys(self, value: List[Union[str, Tuple[str]]]):
-        # the first time out_keys are set, they are marked as ground truth
-        value = unravel_key_list(list(value))
-        if not hasattr(self, "_out_keys"):
-            self._out_keys = value
-        self._out_keys_apparent = value
+        self._out_keys = self._out_keys_apparent = unravel_key_list(list(value))
 
     def select_out_keys(self, *out_keys) -> TensorDictModuleBase:  # noqa: F811
         """Selects the keys that will be found in the output tensordict.
@@ -676,10 +682,9 @@ def select_out_keys(self, *out_keys) -> TensorDictModuleBase:  # noqa: F811
                 ):
                     err_msg += f"Are you passing the keys in a list? Try unpacking as: `{', '.join(out_keys[0])}`"
                 raise ValueError(err_msg)
-        self.register_forward_hook(_OutKeysSelect(out_keys), with_kwargs=True)
-        for hook in self._forward_hooks.values():
-            if isinstance(hook, _OutKeysSelect):
-                hook._init(self)
+        hook = _OutKeysSelect(out_keys)
+        hook._init(self)
+        self.register_forward_hook(hook, with_kwargs=True)
         return self
 
     def reset_out_keys(self):
@@ -1324,6 +1329,19 @@ def __init__(self, td_module: TensorDictModuleBase) -> None:
             for pre_hook in self.td_module._forward_hooks:
                 self.register_forward_hook(self.td_module._forward_hooks[pre_hook])
 
+    @property
+    def out_keys(self):
+        return self.__dict__.get("_out_keys_apparent", self.td_module.out_keys)
+
+    @property
+    def out_keys_source(self):
+        return self.td_module.out_keys_source
+
+    @out_keys.setter
+    def out_keys(self, value: List[Union[str, Tuple[str]]]):
+        self.td_module.out_keys = value
+        self._out_keys_apparent = self.td_module.out_keys
+
     def __getattr__(self, name: str) -> Any:
         if not is_compiling():
             __dict__ = self.__dict__

tensordict/nn/sequence.py

@@ -191,6 +191,15 @@ class TensorDictSequential(TensorDictModule):
     module: nn.ModuleList
     _select_before_return = False
 
+    @property
+    def out_keys(self):
+        return self._out_keys_apparent
+
+    @out_keys.setter
+    def out_keys(self, value: List[NestedKey]):
+        self._out_keys = self._out_keys_apparent = unravel_key_list(list(value))
+        self._complete_out_keys = list(self._out_keys)
+
     @overload
     def __init__(
         self,
@@ -275,7 +284,7 @@ def __init__(
             selected_out_keys = unravel_key_list(selected_out_keys)
             if not all(key in self.out_keys for key in selected_out_keys):
                 raise ValueError("All keys in selected_out_keys must be in out_keys.")
-            self.out_keys = selected_out_keys
+            self._out_keys_apparent = selected_out_keys
         else:
             self._select_before_return = False
 
@@ -337,7 +346,7 @@ def select_out_keys(self, *selected_out_keys) -> TensorDictSequential:
         selected_out_keys = unravel_key_list(selected_out_keys)
         if not all(key in self.out_keys for key in selected_out_keys):
             raise ValueError("All keys in selected_out_keys must be in out_keys.")
-        self.out_keys = selected_out_keys
+        self._out_keys_apparent = selected_out_keys
         return self
 
     def select_subsequence(

test/nn/test_nn.py

@@ -408,6 +408,45 @@ def test_mutable_sequence(self):
         assert "e" in td
         assert "f" in td
 
+    @pytest.mark.parametrize("wrap", [False, True])
+    def test_out_keys_setter(self, wrap):
+        net = nn.Linear(3, 4)
+        td_module = TensorDictModule(net, in_keys=["in"], out_keys=["out"])
+        if wrap:
+            td_module = TensorDictModuleWrapper(td_module)
+
+        def split_output(module, args, output):
+            return output, output.mean(dim=1)
+
+        net.register_forward_hook(split_output)
+        td_module.out_keys = ["out1", "out2"]
+
+        x = torch.randn(3, 3)
+        td_out = td_module(TensorDict({"in": x}, [3]))
+        assert td_module.out_keys == td_module.out_keys_source == ["out1", "out2"]
+        assert "out" not in td_out
+        torch.testing.assert_close(td_out["out2"], td_out["out1"].mean(dim=1))
+
+        out1, out2 = td_module(x)
+        torch.testing.assert_close(out2, out1.mean(dim=1))
+
+        td_module.select_out_keys("out1")
+        td_module.out_keys = ["out1", "out2"]
+        assert "out2" in td_module(TensorDict({"in": x}, [3]))
+
+    def test_select_out_keys_unsupported_property(self):
+        module = ProbabilisticTensorDictModule(
+            in_keys=["loc", "scale"],
+            out_keys=["sample"],
+            distribution_class=Normal,
+            return_log_prob=True,
+        )
+
+        with pytest.raises(RuntimeError, match="does not support select_out_keys"):
+            module.select_out_keys("sample")
+        td = module(TensorDict({"loc": torch.zeros(()), "scale": torch.ones(())}, []))
+        assert set(module.out_keys).issubset(td.keys())
+
     def test_auto_unravel(self):
         tdm = TensorDictModule(
             lambda x: x,
@@ -1271,6 +1310,21 @@ def test_key_exclusion_constructor(self):
         )
         assert set(seq.in_keys) == set(unravel_key_list(("key1", "key2", "key3")))
         assert seq.out_keys == ["key2"]
+        assert seq.out_keys_source == ["foo1", "key1", "key2"]
+
+    def test_out_keys_setter(self):
+        module = TensorDictModule(
+            lambda x: (x, x + 1), in_keys=["in"], out_keys=["out1", "out2"]
+        )
+        seq = TensorDictSequential(module)
+        module.out_keys = ["new_out1", "new_out2"]
+        seq.out_keys = module.out_keys
+
+        seq.select_out_keys("new_out2").reset_out_keys()
+
+        assert seq.out_keys == seq.out_keys_source == ["new_out1", "new_out2"]
+        out1, out2 = seq(torch.zeros(()))
+        torch.testing.assert_close(out2, out1 + 1)
 
     def test_key_exclusion_constructor_exec(self):
         module1 = TensorDictModule(
@@ -2022,6 +2076,7 @@ def test_tdmodule(self, out_d_key, unpack):
             mod2 = mod.select_out_keys(*out_d_key)
             assert mod2 is mod
             assert mod.out_keys == unravel_key_list(out_d_key)
+            assert mod.out_keys_source == ["c", "d", "e"]
             td = mod(TensorDict({"a": torch.zeros(()), "b": torch.ones(())}, []))
             assert "c" not in td.keys()
             assert all(key in td.keys() for key in ["a", "b", "d"])
@@ -2148,6 +2203,7 @@ def test_tdseq(self, out_d_key, unpack):
             mod2 = mod.select_out_keys(*out_d_key)
             assert mod2 is mod
             assert mod.out_keys == unravel_key_list(out_d_key)
+            assert mod.out_keys_source == ["c", "d", "e"]
             td = mod(TensorDict({"a": torch.zeros(()), "b": torch.ones(())}, []))
             assert "c" not in td.keys()
             assert all(key in td.keys() for key in ["a", "b", "d"])