Coverage for transformer_lens/hook_points.py: 91%
202 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1from __future__ import annotations
3"""Hook Points.
5Helpers to access activations in models.
6"""
8from collections.abc import Callable, Sequence
9from dataclasses import dataclass
10from functools import partial
11from typing import (
12 Any,
13 Callable,
14 Literal,
15 Optional,
16 Protocol,
17 Sequence,
18 Union,
19 runtime_checkable,
20)
22import torch
23import torch.nn as nn
24import torch.utils.hooks as hooks
25from torch import Tensor
27from transformer_lens.conversion_utils.conversion_steps.base_tensor_conversion import (
28 BaseTensorConversion,
29)
32@dataclass
33class LensHandle:
34 """Dataclass that holds information about a PyTorch hook."""
36 hook: hooks.RemovableHandle
37 """Reference to the Hook's Removable Handle."""
39 is_permanent: bool = False
40 """Indicates if the Hook is Permanent."""
42 context_level: Optional[int] = None
43 """Context level associated with the hooks context manager for the given hook."""
45 user_hook: Optional[Callable] = None
46 """The original hook callable, before ``add_hook`` wraps it."""
49# Define type aliases
50NamesFilter = Optional[Union[Callable[[str], bool], Sequence[str], str]]
53class _ScaledGradientTensor:
54 """Wrapper around gradient tensors that applies backward_scale to sum operations.
56 This works around a PyTorch bug/behavior where multiplying gradient tensors
57 element-wise in backward hooks gives incorrect sums.
58 """
60 def __init__(self, tensor: Tensor, scale: float):
61 self._tensor = tensor
62 self._scale = scale
64 def sum(self, *args, **kwargs):
65 """Override sum to apply scaling to the result, not the tensor."""
66 result = self._tensor.sum(*args, **kwargs)
67 if isinstance(result, Tensor) and result.numel() == 1:
68 # Scalar result - apply scale
69 return result * self._scale
70 return result
72 def __getattr__(self, name):
73 """Delegate all other attributes to the wrapped tensor."""
74 return getattr(self._tensor, name)
76 def __repr__(self):
77 return f"ScaledGradientTensor({self._tensor}, scale={self._scale})"
80@runtime_checkable
81class _HookFunctionProtocol(Protocol):
82 """Protocol for hook functions."""
84 def __call__(self, tensor: Tensor, *, hook: "HookPoint") -> Union[Any, None]:
85 ...
88HookFunction = _HookFunctionProtocol # Callable[..., _HookFunctionProtocol]
90DeviceType = Optional[torch.device]
91_grad_t = Union[tuple[Tensor, ...], Tensor]
94class _AliasedHookPoint:
95 """
96 A lightweight wrapper that represents a HookPoint with an aliased name.
98 This is used when a hook is registered with multiple names (e.g., in compatibility mode
99 where both canonical and legacy names should trigger the hook). Instead of modifying
100 the original HookPoint's name, we create this wrapper that delegates to the original
101 HookPoint but presents a different name to the user's hook function.
102 """
104 def __init__(self, alias_name: str, target: "HookPoint"):
105 """
106 Create an aliased view of a HookPoint.
108 Args:
109 alias_name: The name to present to the hook function
110 target: The original HookPoint to delegate to
111 """
112 self._alias_name = alias_name
113 self._target = target
115 @property
116 def name(self) -> Optional[str]:
117 """Return the alias name."""
118 return self._alias_name
120 @property
121 def ctx(self) -> dict:
122 """Delegate to the target's context."""
123 return self._target.ctx
125 @property
126 def hook_conversion(self):
127 """Delegate to the target's hook conversion."""
128 return self._target.hook_conversion
130 def layer(self) -> int:
131 """
132 Extract layer index from the alias name.
134 Returns the layer index for hook names like 'blocks.0.attn.hook_pattern' -> 0
135 """
136 if self._alias_name is None:
137 raise ValueError("Name cannot be None")
138 split_name = self._alias_name.split(".")
139 return int(split_name[1])
142class HookPoint(nn.Module):
143 """
144 A helper class to access intermediate activations in a PyTorch model (inspired by Garcon).
146 HookPoint is a dummy module that acts as an identity function by default. By wrapping any
147 intermediate activation in a HookPoint, it provides a convenient way to add PyTorch hooks.
148 """
150 def __init__(self):
151 super().__init__()
152 self.fwd_hooks: list[LensHandle] = []
153 self.bwd_hooks: list[LensHandle] = []
154 self.ctx = {}
156 # A variable giving the hook's name (from the perspective of the root
157 # module) - this is set by the root module at setup.
158 self.name: Optional[str] = None
160 # Hook conversion for input and output transformations
161 self.hook_conversion: Optional[BaseTensorConversion] = None
163 # Backward gradient scale factor (for compatibility between architectures)
164 # This scales the SUM of gradients, not element-wise (to avoid PyTorch bugs)
165 self.backward_scale: float = 1.0
167 def __repr__(self) -> str:
168 bits = [f"name={self.name!r}"] if self.name is not None else []
169 if self.fwd_hooks:
170 bits.append(f"{len(self.fwd_hooks)} fwd")
171 if self.bwd_hooks: 171 ↛ 172line 171 didn't jump to line 172 because the condition on line 171 was never true
172 bits.append(f"{len(self.bwd_hooks)} bwd")
173 return f"HookPoint({', '.join(bits)})" if bits else "HookPoint()"
175 def add_perma_hook(self, hook: HookFunction, dir: Literal["fwd", "bwd"] = "fwd") -> None:
176 self.add_hook(hook, dir=dir, is_permanent=True)
178 def add_hook(
179 self,
180 hook: HookFunction,
181 dir: Literal["fwd", "bwd"] = "fwd",
182 is_permanent: bool = False,
183 level: Optional[int] = None,
184 prepend: bool = False,
185 alias_names: Optional[list[str]] = None,
186 ) -> None:
187 """
188 Hook format is fn(activation, hook_name)
189 Change it into PyTorch hook format (this includes input and output,
190 which are the same for a HookPoint)
191 If prepend is True, add this hook before all other hooks
192 If alias_names is provided, the hook will be called once for each alias name,
193 receiving a temporary HookPoint-like object with that name instead of self
194 (useful for compatibility mode aliases)
195 """
197 def full_hook(
198 module: torch.nn.Module,
199 module_input: Any,
200 module_output: Any,
201 ):
202 if (
203 dir == "bwd"
204 ): # For a backwards hook, module_output is a tuple of (grad,) - I don't know why.
205 module_output = module_output[0]
207 # Apply backward scaling if needed (wrap tensor to scale sum operations)
208 if self.backward_scale != 1.0: 208 ↛ 209line 208 didn't jump to line 209 because the condition on line 208 was never true
209 module_output = _ScaledGradientTensor(module_output, self.backward_scale)
211 # Apply input conversion if hook_conversion exists
212 if self.hook_conversion is not None:
213 module_output = self.hook_conversion.convert(module_output)
215 # Apply the hook for each name (or just once with canonical name)
216 if alias_names is not None:
217 # Call the hook once for each alias name
218 # Create a simple wrapper that acts like a HookPoint but with a different name
219 hook_result = None
220 hook_changed_output = False
221 for alias_name in alias_names:
222 # Create a view of this HookPoint with the alias name
223 hook_with_alias = _AliasedHookPoint(alias_name, self)
224 hook_result = hook(module_output, hook=hook_with_alias) # type: ignore[arg-type]
226 # If the hook modified the output, use that for subsequent calls
227 if hook_result is not None:
228 module_output = hook_result
229 hook_changed_output = True
230 if hook_changed_output: 230 ↛ 237line 230 didn't jump to line 237 because the condition on line 230 was always true
231 hook_result = module_output
232 else:
233 # Call the hook once with the canonical name (self)
234 hook_result = hook(module_output, hook=self)
236 # Apply output reversion if hook_conversion exists and hook returned a value
237 if hook_result is not None and self.hook_conversion is not None:
238 hook_result = self.hook_conversion.revert(hook_result)
240 # For backward hooks, PyTorch expects the return to be a tuple of (grad,)
241 if dir == "bwd" and hook_result is not None:
242 return (
243 hook_result
244 if isinstance(hook_result, tuple) and len(hook_result) == 1
245 else (hook_result,)
246 )
248 return hook_result
250 # annotate the `full_hook` with the string representation of the `hook` function
251 if isinstance(hook, partial):
252 # partial.__repr__() can be extremely slow if arguments contain large objects, which
253 # is common when caching tensors.
254 full_hook.__name__ = f"partial({hook.func.__repr__()},...)"
255 else:
256 full_hook.__name__ = hook.__repr__()
258 if dir == "fwd":
259 pt_handle = self.register_forward_hook(full_hook, prepend=prepend)
260 visible_hooks = self.fwd_hooks
261 elif dir == "bwd": 261 ↛ 282line 261 didn't jump to line 282 because the condition on line 261 was always true
262 # Wrap full_hook's bare Tensor return in tuple for PyTorch's backward API
263 def _bwd_hook_wrapper(
264 module: torch.nn.Module,
265 grad_input: Any,
266 grad_output: Any,
267 ):
268 result = full_hook(module, grad_input, grad_output)
269 if result is None:
270 return None
271 if isinstance(result, tuple): 271 ↛ 273line 271 didn't jump to line 273 because the condition on line 271 was always true
272 return result
273 return (result,)
275 if isinstance(hook, partial):
276 _bwd_hook_wrapper.__name__ = f"partial({hook.func.__repr__()},...)"
277 else:
278 _bwd_hook_wrapper.__name__ = hook.__repr__()
279 pt_handle = self.register_full_backward_hook(_bwd_hook_wrapper, prepend=prepend)
280 visible_hooks = self.bwd_hooks
281 else:
282 raise ValueError(f"Invalid direction {dir}")
284 handle = LensHandle(pt_handle, is_permanent, level, user_hook=hook)
286 if prepend:
287 # we could just pass this as an argument in PyTorch 2.0, but for now we manually do this...
288 visible_hooks.insert(0, handle)
290 else:
291 visible_hooks.append(handle)
293 def has_hooks(
294 self,
295 dir: Literal["fwd", "bwd", "both"] = "both",
296 including_permanent: bool = True,
297 level: Optional[int] = None,
298 ) -> bool:
299 """Check if this HookPoint has any active hooks.
301 Args:
302 dir: Direction of hooks to check ("fwd", "bwd", or "both")
303 including_permanent: Whether to include permanent hooks in the check
304 level: Only check hooks at this context level (None for all levels)
306 Returns:
307 True if any matching hooks are found, False otherwise
308 """
310 def _has_hooks_in_direction(handles: list[LensHandle]) -> bool:
311 for handle in handles:
312 # Check if this hook matches our criteria
313 if not including_permanent and handle.is_permanent:
314 continue
315 if level is not None and handle.context_level != level:
316 continue
317 return True
318 return False
320 if dir == "fwd":
321 return _has_hooks_in_direction(self.fwd_hooks)
322 elif dir == "bwd":
323 return _has_hooks_in_direction(self.bwd_hooks)
324 elif dir == "both": 324 ↛ 329line 324 didn't jump to line 329 because the condition on line 324 was always true
325 return _has_hooks_in_direction(self.fwd_hooks) or _has_hooks_in_direction(
326 self.bwd_hooks
327 )
328 else:
329 raise ValueError(f"Invalid direction {dir}")
331 def remove_hooks(
332 self,
333 dir: Literal["fwd", "bwd", "both"] = "fwd",
334 including_permanent: bool = False,
335 level: Optional[int] = None,
336 ) -> None:
337 def _remove_hooks(handles: list[LensHandle]) -> list[LensHandle]:
338 output_handles = []
339 for handle in handles:
340 if including_permanent:
341 handle.hook.remove()
342 elif (not handle.is_permanent) and (level is None or handle.context_level == level):
343 handle.hook.remove()
344 else:
345 output_handles.append(handle)
346 return output_handles
348 if dir == "fwd" or dir == "both":
349 self.fwd_hooks = _remove_hooks(self.fwd_hooks)
350 if dir == "bwd" or dir == "both":
351 self.bwd_hooks = _remove_hooks(self.bwd_hooks)
352 if dir not in ["fwd", "bwd", "both"]: 352 ↛ 353line 352 didn't jump to line 353 because the condition on line 352 was never true
353 raise ValueError(f"Invalid direction {dir}")
355 def clear_context(self):
356 del self.ctx
357 self.ctx = {}
359 def enable_reshape(
360 self,
361 hook_conversion: Optional[BaseTensorConversion] = None,
362 ) -> None:
363 """
364 Enable reshape functionality for this hook point using a BaseTensorConversion.
366 Args:
367 hook_conversion: BaseTensorConversion instance to handle input/output transformations.
368 The convert() method will be used for input transformation,
369 and the revert() method will be used for output transformation.
370 """
371 self.hook_conversion = hook_conversion
373 def forward(self, x: Tensor) -> Tensor:
374 return x
376 def layer(self):
377 # Returns the layer index if the name has the form 'blocks.{layer}.{...}'
378 # Helper function that's mainly useful on HookedTransformer
379 # If it doesn't have this form, raises an error -
380 if self.name is None: 380 ↛ 381line 380 didn't jump to line 381 because the condition on line 380 was never true
381 raise ValueError("Name cannot be None")
382 split_name = self.name.split(".")
383 return int(split_name[1])
386# %%
387class HookIntrospectionMixin:
388 """``list_hooks()`` mixins for any class exposing a ``hook_dict``.
390 Accessed via ``getattr`` so subclasses can provide ``hook_dict`` as either
391 an instance attribute (``HookedRootModule``) or a ``@property`` (``TransformerBridge``).
392 """
394 def list_hooks(
395 self,
396 name_filter: NamesFilter = None,
397 dir: Literal["fwd", "bwd", "both"] = "both",
398 including_permanent: bool = True,
399 ) -> dict[str, list[LensHandle]]:
400 """Return attached hooks grouped by HookPoint name; empty HookPoints are omitted.
402 Args:
403 name_filter: A hook name, list of names, or predicate. ``None`` matches all.
404 dir: Restrict to forward, backward, or both directions.
405 including_permanent: If False, drop permanent hooks from the result.
406 """
407 if name_filter is None:
408 matches: Callable[[str], bool] = lambda _: True
409 elif callable(name_filter):
410 matches = name_filter
411 elif isinstance(name_filter, str):
412 target = name_filter
413 matches = lambda n: n == target
414 else:
415 allowed = set(name_filter)
416 matches = lambda n: n in allowed
418 out: dict[str, list[LensHandle]] = {}
419 hook_dict: dict[str, HookPoint] = getattr(self, "hook_dict")
420 for name, hp in hook_dict.items():
421 if not matches(name):
422 continue
423 handles: list[LensHandle] = []
424 if dir in ("fwd", "both"):
425 handles.extend(hp.fwd_hooks)
426 if dir in ("bwd", "both"):
427 handles.extend(hp.bwd_hooks)
428 if not including_permanent:
429 handles = [h for h in handles if not h.is_permanent]
430 if handles:
431 out[name] = handles
432 return out
435# HookedRootModule moved to transformer_lens.HookedRootModule (3.0). Import it from
436# its dedicated module. Importing from here is deprecated and will trigger a warning.
437def __getattr__(name: str):
438 if name == "HookedRootModule":
439 import warnings
441 from transformer_lens.HookedRootModule import HookedRootModule
443 warnings.warn(
444 "Importing HookedRootModule from transformer_lens.hook_points is "
445 "deprecated and will be removed in TransformerLens 4.0. Import it from "
446 "transformer_lens (preferred) or transformer_lens.HookedRootModule instead.",
447 DeprecationWarning,
448 stacklevel=2,
449 )
450 return HookedRootModule
451 raise AttributeError(f"module {__name__!r} has no attribute {name!r}")