From 31aa33c639b827819b0cf9ac83c161e2c9bacaf6 Mon Sep 17 00:00:00 2001 From: Bryce Meyer Date: Wed, 15 Oct 2025 23:18:18 +0200 Subject: [PATCH 1/2] setup new hooks properly --- transformer_lens/hook_points.py | 35 +++++++++++++++++++++++-- transformer_lens/model_bridge/bridge.py | 32 ++++++++++++++++++++-- 2 files changed, 63 insertions(+), 4 deletions(-) diff --git a/transformer_lens/hook_points.py b/transformer_lens/hook_points.py index 698cf0a27..d5e8a103e 100644 --- a/transformer_lens/hook_points.py +++ b/transformer_lens/hook_points.py @@ -98,12 +98,16 @@ def add_hook( is_permanent: bool = False, level: Optional[int] = None, prepend: bool = False, + alias_names: Optional[list[str]] = None, ) -> None: """ Hook format is fn(activation, hook_name) Change it into PyTorch hook format (this includes input and output, which are the same for a HookPoint) If prepend is True, add this hook before all other hooks + If alias_names is provided, the hook will be called once for each alias name, + receiving a temporary HookPoint-like object with that name instead of self + (useful for compatibility mode aliases) """ def full_hook( @@ -120,8 +124,35 @@ def full_hook( if self.hook_conversion is not None: module_output = self.hook_conversion.convert(module_output) - # Apply the hook - hook_result = hook(module_output, hook=self) + # Determine what names to use for hook calls + if alias_names is not None: + # Call the hook once for each alias name + names_to_use = alias_names + else: + # Call the hook once with the canonical name + names_to_use = [None] + + # Apply the hook for each name + hook_result = None + for name in names_to_use: + if name is not None: + # Create a lightweight object with the alias name + class _NamedHook: + def __init__(self, name: str, target: "HookPoint"): + self.name = name + self.ctx = target.ctx + self.hook_conversion = target.hook_conversion + + hook_param = _NamedHook(name, self) + else: + hook_param = self + + # Apply the hook + hook_result = hook(module_output, hook=hook_param) + + # If the hook modified the output, use that for subsequent calls + if hook_result is not None: + module_output = hook_result # Apply output reversion if hook_conversion exists and hook returned a value if hook_result is not None and self.hook_conversion is not None: diff --git a/transformer_lens/model_bridge/bridge.py b/transformer_lens/model_bridge/bridge.py index 8ead4efd5..28a89c332 100644 --- a/transformer_lens/model_bridge/bridge.py +++ b/transformer_lens/model_bridge/bridge.py @@ -968,7 +968,21 @@ def _run_with_hooks_ported( def add_hook_to_point( hook_point: HookPoint, hook_fn: Callable, name: str, dir: str = "fwd" ): - hook_point.add_hook(hook_fn, dir=dir) # type: ignore[arg-type] + # In compatibility mode, collect all names for this hook point + # (canonical name + any aliases that map to it) + if self.compatibility_mode: + alias_names_list = [] + + # Add the canonical name first + alias_names_list.append(hook_point.name) + + # Add any alias names that differ from the canonical name + if name != hook_point.name: + alias_names_list.append(name) + + hook_point.add_hook(hook_fn, dir=dir, alias_names=alias_names_list) # type: ignore[arg-type] + else: + hook_point.add_hook(hook_fn, dir=dir) # type: ignore[arg-type] added_hooks.append((hook_point, name)) try: @@ -4244,7 +4258,21 @@ def run_with_hooks( def add_hook_to_point( hook_point: HookPoint, hook_fn: Callable, name: str, dir: Literal["fwd", "bwd"] = "fwd" ): - hook_point.add_hook(hook_fn, dir=dir) + # In compatibility mode, collect all names for this hook point + # (canonical name + any aliases that map to it) + if self.compatibility_mode: + alias_names_list = [] + + # Add the canonical name first + alias_names_list.append(hook_point.name) + + # Add any alias names that differ from the canonical name + if name != hook_point.name: + alias_names_list.append(name) + + hook_point.add_hook(hook_fn, dir=dir, alias_names=alias_names_list) + else: + hook_point.add_hook(hook_fn, dir=dir) added_hooks.append((hook_point, name)) # Add stop_at_layer hook if specified From f9831eca867b07f8879aca7f87a2fe4fc327fdbd Mon Sep 17 00:00:00 2001 From: Bryce Meyer Date: Wed, 15 Oct 2025 23:33:59 +0200 Subject: [PATCH 2/2] fixed type checks --- transformer_lens/hook_points.py | 43 +++++++++++-------------- transformer_lens/model_bridge/bridge.py | 10 +++--- 2 files changed, 24 insertions(+), 29 deletions(-) diff --git a/transformer_lens/hook_points.py b/transformer_lens/hook_points.py index d5e8a103e..f7956eb6a 100644 --- a/transformer_lens/hook_points.py +++ b/transformer_lens/hook_points.py @@ -124,35 +124,28 @@ def full_hook( if self.hook_conversion is not None: module_output = self.hook_conversion.convert(module_output) - # Determine what names to use for hook calls + # Apply the hook for each name (or just once with canonical name) if alias_names is not None: # Call the hook once for each alias name - names_to_use = alias_names - else: - # Call the hook once with the canonical name - names_to_use = [None] - - # Apply the hook for each name - hook_result = None - for name in names_to_use: - if name is not None: - # Create a lightweight object with the alias name - class _NamedHook: - def __init__(self, name: str, target: "HookPoint"): - self.name = name - self.ctx = target.ctx - self.hook_conversion = target.hook_conversion - + # Define _NamedHook class to match HookPoint protocol + class _NamedHook: + def __init__(self, name: str, target: "HookPoint"): + self.name = name + self.ctx = target.ctx + self.hook_conversion = target.hook_conversion + + hook_result = None + for name in alias_names: hook_param = _NamedHook(name, self) - else: - hook_param = self + # Apply the hook + hook_result = hook(module_output, hook=hook_param) # type: ignore[arg-type] - # Apply the hook - hook_result = hook(module_output, hook=hook_param) - - # If the hook modified the output, use that for subsequent calls - if hook_result is not None: - module_output = hook_result + # If the hook modified the output, use that for subsequent calls + if hook_result is not None: + module_output = hook_result + else: + # Call the hook once with the canonical name (self) + hook_result = hook(module_output, hook=self) # Apply output reversion if hook_conversion exists and hook returned a value if hook_result is not None and self.hook_conversion is not None: diff --git a/transformer_lens/model_bridge/bridge.py b/transformer_lens/model_bridge/bridge.py index 28a89c332..f5aa7d4b4 100644 --- a/transformer_lens/model_bridge/bridge.py +++ b/transformer_lens/model_bridge/bridge.py @@ -971,10 +971,11 @@ def add_hook_to_point( # In compatibility mode, collect all names for this hook point # (canonical name + any aliases that map to it) if self.compatibility_mode: - alias_names_list = [] + alias_names_list: list[str] = [] # Add the canonical name first - alias_names_list.append(hook_point.name) + if hook_point.name is not None: + alias_names_list.append(hook_point.name) # Add any alias names that differ from the canonical name if name != hook_point.name: @@ -4261,10 +4262,11 @@ def add_hook_to_point( # In compatibility mode, collect all names for this hook point # (canonical name + any aliases that map to it) if self.compatibility_mode: - alias_names_list = [] + alias_names_list: list[str] = [] # Add the canonical name first - alias_names_list.append(hook_point.name) + if hook_point.name is not None: + alias_names_list.append(hook_point.name) # Add any alias names that differ from the canonical name if name != hook_point.name: