made sure to check for nested hooks (#1035)

* made sure to check for nested hooks

* removed extra check

* Skip original_model in hook scanning

* Remove extra traversal through general components submodules

* Remove adding aliases to hook registry

* Fix typing error

* Fix typing error

* Remove constant copying of hook dictionary to save memory

---------

Co-authored-by: degenfabian <[email protected]>

Author: bryce13950

transformer_lens/model_bridge/bridge.py

@@ -131,28 +131,6 @@ def _initialize_hook_registry(self) -> None:
         # Scan existing components for hooks
         self._scan_existing_hooks(self, "")
 
-        # Add bridge aliases if compatibility mode is enabled
-        if self.compatibility_mode and self.hook_aliases:
-            for alias_name, target in self.hook_aliases.items():
-                # Use the existing alias system to resolve the target hook
-                # Convert to Dict[str, str] for resolve_alias if target_name is a list
-                if isinstance(target, list):
-                    # For list targets, try each one until one works
-                    for single_target in target:
-                        try:
-                            target_hook = resolve_alias(
-                                self, alias_name, {alias_name: single_target}
-                            )
-                            if target_hook is not None:
-                                self._hook_registry[alias_name] = target_hook
-                                break
-                        except AttributeError:
-                            continue
-                else:
-                    target_hook = resolve_alias(self, alias_name, {alias_name: target})
-                    if target_hook is not None:
-                        self._hook_registry[alias_name] = target_hook
-
         self._hook_registry_initialized = True
 
     def _scan_existing_hooks(self, module: nn.Module, prefix: str = "") -> None:
@@ -180,41 +158,39 @@ def scan_module(mod: nn.Module, path: str = "") -> None:
                 except Exception:
                     # If get_hooks() fails, fall through to the else block
                     pass
-            else:
-                # Fall back to scanning attributes for non-GeneralizedComponent modules
-                for attr_name in dir(mod):
-                    if attr_name.startswith("_"):
-                        continue
-                    if attr_name == "original_component":
-                        continue
 
-                    try:
-                        attr = getattr(mod, attr_name)
-                    except Exception:
-                        continue
-
-                    name = f"{path}.{attr_name}" if path else attr_name
-
-                    if isinstance(attr, HookPoint):
-                        attr.name = name
-                        self._hook_registry[name] = attr
-                    elif isinstance(attr, HookPointWrapper):
-                        hook_in_name = f"{name}.hook_in"
-                        hook_out_name = f"{name}.hook_out"
-                        attr.hook_in.name = hook_in_name
-                        attr.hook_out.name = hook_out_name
-                        self._hook_registry[hook_in_name] = attr.hook_in
-                        self._hook_registry[hook_out_name] = attr.hook_out
-                    elif isinstance(attr, nn.Module) and attr is not mod:
-                        scan_module(attr, name)
-                    elif isinstance(attr, (list, tuple)):
-                        for i, item in enumerate(attr):
-                            if isinstance(item, nn.Module):
-                                scan_module(item, f"{name}[{i}]")
+            # Always scan attributes for additional hooks and submodules
+            for attr_name in dir(mod):
+                if attr_name.startswith("_"):
+                    continue
+                if attr_name == "original_component" or "original_model":
+                    continue
+
+                try:
+                    attr = getattr(mod, attr_name)
+                except Exception:
+                    continue
+
+                name = f"{path}.{attr_name}" if path else attr_name
+
+                if isinstance(attr, HookPoint):
+                    attr.name = name
+                    self._hook_registry[name] = attr
+                elif isinstance(attr, HookPointWrapper):
+                    hook_in_name = f"{name}.hook_in"
+                    hook_out_name = f"{name}.hook_out"
+                    attr.hook_in.name = hook_in_name
+                    attr.hook_out.name = hook_out_name
+                    self._hook_registry[hook_in_name] = attr.hook_in
+                    self._hook_registry[hook_out_name] = attr.hook_out
 
             # Check named children
             for child_name, child_module in mod.named_children():
-                if child_name == "original_component" or child_name == "_original_component":
+                if (
+                    child_name == "original_component"
+                    or child_name == "_original_component"
+                    or child_name == "original_model"
+                ):
                     continue
                 child_path = f"{path}.{child_name}" if path else child_name
                 scan_module(child_module, child_path)
@@ -225,23 +201,7 @@ def scan_module(mod: nn.Module, path: str = "") -> None:
     def hook_dict(self) -> dict[str, HookPoint]:
         """Get all HookPoint objects in the model for compatibility with HookedTransformer."""
         # Start with the current registry
-        hooks = self._hook_registry.copy()
-
-        # Add aliases if compatibility mode is enabled
-        if self.compatibility_mode:
-            for alias_name, target in self.hook_aliases.items():
-                # Handle both string and list target names
-                if isinstance(target, list):
-                    # For list targets, find the first one that exists in hooks
-                    for single_target in target:
-                        if single_target in hooks:
-                            hooks[alias_name] = hooks[single_target]
-                            break
-                else:
-                    if target in hooks:
-                        hooks[alias_name] = hooks[target]
-
-        return hooks
+        return self._hook_registry.copy()
 
     def _discover_hooks(self) -> dict[str, HookPoint]:
         """Get all HookPoint objects from the registry (deprecated, use hook_dict)."""
@@ -262,17 +222,10 @@ def __getattr__(self, name: str) -> Any:
             return self.__dict__[name]
 
         # Check if this is a hook alias when compatibility mode is enabled
-        if self.compatibility_mode and name in self.hook_aliases:
-            target = self.hook_aliases[name]
-            # Handle both string and list target names
-            if isinstance(target, list):
-                # For list targets, find the first one that exists in the registry
-                for single_target in target:
-                    if single_target in self._hook_registry:
-                        return self._hook_registry[single_target]
-            else:
-                if target in self._hook_registry:
-                    return self._hook_registry[target]
+        if self.compatibility_mode:
+            resolved_hook = resolve_alias(self, name, self.hook_aliases)
+            if resolved_hook is not None:
+                return resolved_hook
 
         return super().__getattr__(name)
 
@@ -982,7 +935,7 @@ def cache_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
             return cache_hook
 
         # Use cached hooks instead of re-discovering them
-        hook_dict = self.hook_dict
+        hook_dict = self._hook_registry
 
         # Filter hooks based on names_filter
         for hook_name, hook in hook_dict.items():
@@ -1026,7 +979,7 @@ def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
 
                     # Add hook to the output of the last layer to be processed
                     block_hook_name = f"blocks.{last_layer_to_process}.hook_out"
-                    hook_dict = self.hook_dict
+                    hook_dict = self._hook_registry
                     if block_hook_name in hook_dict:
                         hook_dict[block_hook_name].add_hook(stop_hook)
                         hooks.append((hook_dict[block_hook_name], block_hook_name))
@@ -1167,7 +1120,7 @@ def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
 
                 # Add hook to the output of the last layer to be processed
                 block_hook_name = f"blocks.{last_layer_to_process}.hook_out"
-                hook_dict = self.hook_dict
+                hook_dict = self._hook_registry
                 if block_hook_name in hook_dict:
                     add_hook_to_point(hook_dict[block_hook_name], stop_hook, block_hook_name)
 
@@ -1197,7 +1150,7 @@ def wrapped_hook_fn(tensor, hook):
 
                 if isinstance(hook_name_or_filter, str):
                     # Direct hook name - check for aliases first
-                    hook_dict = self.hook_dict
+                    hook_dict = self._hook_registry
                     actual_hook_name = hook_name_or_filter
 
                     # If this is an alias, resolve it to the actual hook name
@@ -1208,7 +1161,7 @@ def wrapped_hook_fn(tensor, hook):
                         add_hook_to_point(hook_dict[actual_hook_name], hook_fn, actual_hook_name)
                 else:
                     # Filter function
-                    hook_dict = self.hook_dict
+                    hook_dict = self._hook_registry
                     for name, hook_point in hook_dict.items():
                         if hook_name_or_filter(name):
                             add_hook_to_point(hook_point, hook_fn, name)

transformer_lens/model_bridge/generalized_components/base.py

@@ -69,10 +69,6 @@ def __init__(
             self.hook_in.hook_conversion = self.conversion_rule
             self.hook_out.hook_conversion = self.conversion_rule
 
-        # Register the standard hooks
-        self._register_hook("hook_in", self.hook_in)
-        self._register_hook("hook_out", self.hook_out)
-
     def _register_hook(self, name: str, hook: HookPoint) -> None:
         """Register a hook in the component's hook registry."""
         # Set the name on the HookPoint
@@ -82,17 +78,20 @@ def _register_hook(self, name: str, hook: HookPoint) -> None:
 
     def get_hooks(self) -> Dict[str, HookPoint]:
         """Get all hooks registered in this component."""
-        hooks = self._hook_registry.copy()
 
         # Add aliases if compatibility mode is enabled
         if self.compatibility_mode and self.hook_aliases:
+            # Only copy hook registry if compatibility mode is enabled to save memory
+            hooks = self._hook_registry.copy()
+
             for alias_name, target_name in self.hook_aliases.items():
                 # Use the existing alias system to resolve the target hook
                 target_hook = resolve_alias(self, alias_name, self.hook_aliases)
                 if target_hook is not None:
                     hooks[alias_name] = target_hook
-
-        return hooks
+            return hooks
+        else:
+            return self._hook_registry
 
     def _is_getattr_called_internally(self) -> bool:
         """This function checks if the __getattr__ method was being called internally

transformer_lens/utilities/aliases.py

@@ -1,13 +1,13 @@
 """Utilities for handling hook aliases in the bridge system."""
 
 import warnings
-from typing import Any, Dict, Optional, Set
+from typing import Any, Dict, List, Optional, Set, Union
 
 
 def resolve_alias(
     target_object: Any,
     requested_name: str,
-    aliases: Dict[str, str],
+    aliases: Dict[str, str] | Dict[str, Union[str, List[str]]],
 ) -> Optional[Any]:
     """Resolve a hook alias to the actual hook object.
 
@@ -76,7 +76,7 @@ def _resolve_single_target(target_name: str) -> Any:
 
 
 def _collect_aliases_from_module(
-    module: Any, path: str, aliases: Dict[str, str], visited: Set[int]
+    module: Any, path: str, aliases: Dict[str, str], visited: Set[int] = set()
 ) -> None:
     """Helper function to collect all aliases from a single module.
     Args: