From a2ff15e1453b398bc843d90b5e89a8b9b8cb79fa Mon Sep 17 00:00:00 2001
From: kigland <shuaizhicheng336@gmail.com>
Date: Tue, 25 Aug 2026 09:09:38 +0800
Subject: [PATCH] Preserve alias hook replacements across no-op callbacks

---
 tests/unit/test_hook_points.py  | 78 +++++++++++++++++++++++++++++++++
 transformer_lens/hook_points.py |  4 ++
 2 files changed, 82 insertions(+)

diff --git a/tests/unit/test_hook_points.py b/tests/unit/test_hook_points.py
index eabd3c3eb..6637f83b0 100644
--- a/tests/unit/test_hook_points.py
+++ b/tests/unit/test_hook_points.py
@@ -211,6 +211,84 @@ def test_hook(activation, hook):
     assert torch.equal(result, test_input)
 
 
+def test_alias_hook_preserves_earlier_replacement_when_later_alias_returns_none():
+    """A later no-op alias must not discard an earlier alias replacement."""
+    import torch
+
+    hook_point = HookPoint()
+    seen = []
+
+    def selective_hook(activation, hook):
+        seen.append((hook.name, activation.item()))
+        if hook.name == "replace":
+            return activation + 10
+        return None
+
+    hook_point.add_hook(selective_hook, alias_names=["replace", "observe"])
+
+    result = hook_point(torch.tensor(1.0))
+
+    assert seen == [("replace", 1.0), ("observe", 11.0)]
+    assert result.item() == 11.0
+
+
+def test_alias_hook_reverts_conversion_after_earlier_replacement():
+    """Alias replacement remains eligible for output conversion reversion."""
+    import torch
+
+    from transformer_lens.conversion_utils.conversion_steps.base_tensor_conversion import (
+        BaseTensorConversion,
+    )
+
+    class ScaleThenShift(BaseTensorConversion):
+        def handle_conversion(self, input_value, *full_context):
+            return input_value * 2
+
+        def revert(self, input_value, *full_context):
+            return input_value + 100
+
+    hook_point = HookPoint()
+    hook_point.enable_reshape(ScaleThenShift())
+
+    def selective_hook(activation, hook):
+        if hook.name == "replace":
+            return activation + 10
+        return None
+
+    hook_point.add_hook(selective_hook, alias_names=["replace", "observe"])
+
+    result = hook_point(torch.tensor(1.0))
+
+    assert result.item() == 112.0
+
+
+def test_backward_alias_hook_preserves_earlier_gradient_replacement():
+    """A later no-op alias must not discard an earlier gradient replacement."""
+    import torch
+
+    hook_point = HookPoint()
+    seen = []
+
+    def selective_hook(gradient, hook):
+        seen.append((hook.name, gradient.item()))
+        if hook.name == "replace":
+            return gradient + 10
+        return None
+
+    hook_point.add_hook(
+        selective_hook,
+        dir="bwd",
+        alias_names=["replace", "observe"],
+    )
+    input_value = torch.tensor(2.0, requires_grad=True)
+
+    (hook_point(input_value) * 3).backward()
+
+    assert seen == [("replace", 3.0), ("observe", 13.0)]
+    assert input_value.grad is not None
+    assert input_value.grad.item() == 13.0
+
+
 class TestHookPointHasHooks:
     """Comprehensive test suite for HookPoint.has_hooks method."""
 
diff --git a/transformer_lens/hook_points.py b/transformer_lens/hook_points.py
index a0df3c887..529d09762 100644
--- a/transformer_lens/hook_points.py
+++ b/transformer_lens/hook_points.py
@@ -217,6 +217,7 @@ def full_hook(
                 # Call the hook once for each alias name
                 # Create a simple wrapper that acts like a HookPoint but with a different name
                 hook_result = None
+                hook_changed_output = False
                 for alias_name in alias_names:
                     # Create a view of this HookPoint with the alias name
                     hook_with_alias = _AliasedHookPoint(alias_name, self)
@@ -225,6 +226,9 @@ def full_hook(
                     # If the hook modified the output, use that for subsequent calls
                     if hook_result is not None:
                         module_output = hook_result
+                        hook_changed_output = True
+                if hook_changed_output:
+                    hook_result = module_output
             else:
                 # Call the hook once with the canonical name (self)
                 hook_result = hook(module_output, hook=self)