Preserve method output contracts with current TensorDict

Author: Xmaster6y

src/tdhook/modules.py

@@ -1,6 +1,6 @@
 from tensordict.nn import TensorDictModuleWrapper, TensorDictModuleBase, TensorDictSequential
 from tensordict import NonTensorData, TensorDict, TensorDictBase
-from tensordict.utils import NestedKey
+from tensordict.utils import NestedKey, unravel_key_list
 from typing import Callable, Optional, TYPE_CHECKING, List
 import torch
 from textwrap import indent
@@ -271,6 +271,19 @@ def __init__(
         self._hooking_context = hooking_context
         self._relative_path = relative_path
 
+    @TensorDictModuleWrapper.out_keys.setter
+    def out_keys(self, value: List[NestedKey]):
+        # Finalization can publish extra outputs that the wrapped model does not
+        # return. Keep their declaration and native key selection on this wrapper.
+        keys = unravel_key_list(list(value))
+        if "_out_keys" not in self.__dict__:
+            self._out_keys = keys
+        self._out_keys_apparent = keys
+
+    @property
+    def out_keys_source(self):
+        return self.__dict__.get("_out_keys", self.td_module.out_keys_source)
+
     @property
     def hook_root(self) -> TensorDictModuleBase:
         """Return the caller-owned module against which hook paths resolve."""

tests/latent/test_activation_caching.py

@@ -11,6 +11,7 @@
 from tdhook.modules import get_best_device
 from tdhook.runtime import HookProgram, HookSpec
 from tdhook.targets import Target
+from tdhook.workflow import Workflow
 
 
 class TestActivationCaching:
@@ -48,6 +49,38 @@ def test_tensordict_execution_publishes_a_native_cache_output(self, default_test
         assert hooked_module.out_keys == ["output", ("activations", "cache")]
         assert result["activations", "cache"]["linear2"].shape == (2, 20)
 
+    @pytest.mark.parametrize("run_in_workflow", [False, True])
+    def test_cache_publication_preserves_the_callers_model_contract(self, run_in_workflow):
+        raw_model = torch.nn.Sequential(torch.nn.Linear(3, 4))
+        input_key = ("inputs", "value")
+        output_key = ("predictions", "value")
+        cache_key = ("activations", "hidden")
+        model = TensorDictModule(raw_model, in_keys=[input_key], out_keys=[output_key])
+        method = ActivationCaching(r"module\.0$", cache_key=cache_key)
+        inputs = torch.randn(2, 3)
+        expected = raw_model(inputs)
+
+        for _ in range(2):
+            data = TensorDict({input_key: inputs}, batch_size=[2])
+            if run_in_workflow:
+                result = Workflow(method)(model, data)
+            else:
+                context = method.prepare(model)
+                assert model.out_keys == [output_key]
+                assert context.module.out_keys == [output_key, cache_key]
+                with context as hooked_module:
+                    result = hooked_module(data)
+
+            assert model.in_keys == [input_key]
+            assert model.out_keys == [output_key]
+            assert model.out_keys_source == [output_key]
+            torch.testing.assert_close(result[output_key], expected)
+            torch.testing.assert_close(result[cache_key]["module.0"], expected)
+
+        plain_result = model(TensorDict({input_key: inputs}, batch_size=[2]))
+        torch.testing.assert_close(plain_result[output_key], expected)
+        assert cache_key not in plain_result.keys(include_nested=True)
+
     def test_target_selection_is_cached_and_reported(self, default_test_model):
         target = Target("linear2", "activation", -1, (0, 2))
 

tests/test_modules.py

@@ -111,3 +111,33 @@ def test_bound_module_is_context_owned_and_finalizes_results(default_test_model)
 
     with pytest.raises(RuntimeError, match="called in context"):
         prepared(TensorDict({"input": torch.ones(2, 10)}, batch_size=[2]))
+
+
+def test_method_outputs_can_be_selected_and_reset_without_changing_model_outputs():
+    class PublishingModule(HookedModule):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self.out_keys = [*self.out_keys, ("metrics", "sum")]
+
+        def finalize_tensordict(self, data):
+            return data.set(("metrics", "sum"), data["output"].sum(-1))
+
+    class PublishingMethod(HookingContextFactory):
+        _hooked_module_class = PublishingModule
+
+    model = TensorDictModule(torch.nn.Identity(), in_keys=["input"], out_keys=["output"])
+    with PublishingMethod().prepare(model) as prepared:
+        assert prepared.out_keys_source == ["output", ("metrics", "sum")]
+        prepared.select_out_keys(("metrics", "sum"))
+        selected = prepared(TensorDict({"input": torch.ones(2, 3)}, batch_size=[2]))
+        assert prepared.out_keys == [("metrics", "sum")]
+        assert "output" not in selected
+        torch.testing.assert_close(selected["metrics", "sum"], torch.full((2,), 3.0))
+
+        prepared.reset_out_keys()
+        assert prepared.out_keys == ["output", ("metrics", "sum")]
+        restored = prepared(TensorDict({"input": torch.ones(2, 3)}, batch_size=[2]))
+        torch.testing.assert_close(restored["output"], torch.ones(2, 3))
+
+    assert model.out_keys == ["output"]
+    assert model.out_keys_source == ["output"]