moved enable hook functionality to separate functions and tested new functions (#613)

Author: bryce13950

tests/unit/test_hooked_root_module.py

@@ -0,0 +1,76 @@
+from unittest.mock import Mock
+
+from transformer_lens.hook_points import HookedRootModule
+
+MODEL_NAME = "solu-2l"
+
+
+def test_enable_hook_with_name():
+    model = HookedRootModule()
+    model.mod_dict = {"linear": Mock()}
+    model.context_level = 5
+
+    hook = lambda x: False
+    dir = "fwd"
+
+    model._enable_hook_with_name("linear", hook=hook, dir=dir)
+
+    model.mod_dict["linear"].add_hook.assert_called_with(hook, dir="fwd", level=5)
+
+
+def test_enable_hooks_for_points():
+    model = HookedRootModule()
+    model.mod_dict = {}
+    model.context_level = 5
+
+    hook_points = {
+        "linear": Mock(),
+        "attn": Mock(),
+    }
+
+    enabled = lambda x: x == "attn"
+
+    hook = lambda x: False
+    dir = "bwd"
+
+    print(hook_points.items())
+    model._enable_hooks_for_points(
+        hook_points=hook_points.items(), enabled=enabled, hook=hook, dir=dir
+    )
+
+    hook_points["attn"].add_hook.assert_called_with(hook, dir="bwd", level=5)
+    hook_points["linear"].add_hook.assert_not_called()
+
+
+def test_enable_hook_with_string_param():
+    model = HookedRootModule()
+    model.mod_dict = {"linear": Mock()}
+    model.context_level = 5
+
+    hook = lambda x: False
+    dir = "fwd"
+
+    model._enable_hook("linear", hook=hook, dir=dir)
+
+    model.mod_dict["linear"].add_hook.assert_called_with(hook, dir="fwd", level=5)
+
+
+def test_enable_hook_with_callable_param():
+    model = HookedRootModule()
+    model.mod_dict = {"linear": Mock()}
+    model.hook_dict = {
+        "linear": Mock(),
+        "attn": Mock(),
+    }
+    model.context_level = 5
+
+    enabled = lambda x: x == "attn"
+
+    hook = lambda x: False
+    dir = "fwd"
+
+    model._enable_hook(enabled, hook=hook, dir=dir)
+
+    model.mod_dict["linear"].add_hook.assert_not_called()
+    model.hook_dict["attn"].add_hook.assert_called_with(hook, dir="fwd", level=5)
+    model.hook_dict["linear"].add_hook.assert_not_called()

transformer_lens/hook_points.py

@@ -7,7 +7,18 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
 from functools import partial
-from typing import Callable, Dict, List, Literal, Optional, Sequence, Tuple, Union, cast
+from typing import (
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Literal,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+    cast,
+)
 
 import torch.nn as nn
 import torch.utils.hooks as hooks
@@ -267,6 +278,50 @@ def add_hook(
     def add_perma_hook(self, name, hook, dir="fwd") -> None:
         self.add_hook(name, hook, dir=dir, is_permanent=True)
 
+    def _enable_hook_with_name(self, name: str, hook: Callable, dir: Literal["fwd", "bwd"]):
+        """This function takes a key for the mod_dict and enables the related hook for that module
+
+        Args:
+            name (str): The module name
+            hook (Callable): The hook to add
+            dir (Literal["fwd", "bwd"]): The direction for the hook
+        """
+        self.mod_dict[name].add_hook(hook, dir=dir, level=self.context_level)
+
+    def _enable_hooks_for_points(
+        self,
+        hook_points: Iterable[Tuple[str, HookPoint]],
+        enabled: Callable,
+        hook: Callable,
+        dir: Literal["fwd", "bwd"],
+    ):
+        """Enables hooks for a list of points
+
+        Args:
+            hook_points (Dict[str, HookPoint]): The hook points
+            enabled (Callable): _description_
+            hook (Callable): _description_
+            dir (Literal["fwd", "bwd"]): _description_
+        """
+        for hook_name, hook_point in hook_points:
+            if enabled(hook_name):
+                hook_point.add_hook(hook, dir=dir, level=self.context_level)
+
+    def _enable_hook(self, name: Union[str, Callable], hook: Callable, dir: Literal["fwd", "bwd"]):
+        """Enables an individual hook on a hook point
+
+        Args:
+            name (str): The name of the hook
+            hook (Callable): The actual hook
+            dir (Literal["fwd", "bwd"], optional): The direction of the hook. Defaults to "fwd".
+        """
+        if isinstance(name, str):
+            self._enable_hook_with_name(name=name, hook=hook, dir=dir)
+        else:
+            self._enable_hooks_for_points(
+                hook_points=self.hook_dict.items(), enabled=name, hook=hook, dir=dir
+            )
+
     @contextmanager
     def hooks(
         self,
@@ -296,21 +351,9 @@ def hooks(
             self.context_level += 1
 
             for name, hook in fwd_hooks:
-                if isinstance(name, str):
-                    self.mod_dict[name].add_hook(hook, dir="fwd", level=self.context_level)
-                else:
-                    # Otherwise, name is a Boolean function on names
-                    for hook_name, hp in self.hook_dict.items():
-                        if name(hook_name):
-                            hp.add_hook(hook, dir="fwd", level=self.context_level)
+                self._enable_hook(name=name, hook=hook, dir="fwd")
             for name, hook in bwd_hooks:
-                if isinstance(name, str):
-                    self.mod_dict[name].add_hook(hook, dir="bwd", level=self.context_level)
-                else:
-                    # Otherwise, name is a Boolean function on names
-                    for hook_name, hp in self.hook_dict:  # type: ignore
-                        if name(hook_name):
-                            hp.add_hook(hook, dir="bwd", level=self.context_level)
+                self._enable_hook(name=name, hook=hook, dir="bwd")
             yield self
         finally:
             if reset_hooks_end: