Coverage for transformer_lens/hook_points.py: 86%
197 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +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."""
49def _handle_matches(handle: LensHandle, including_permanent: bool, level: Optional[int]) -> bool:
50 """Whether a handle is in scope for a ``has_hooks``/``remove_hooks``/``list_hooks`` call."""
51 if not including_permanent and handle.is_permanent:
52 return False
53 return level is None or handle.context_level == level
56# Define type aliases
57NamesFilter = Optional[Union[Callable[[str], bool], Sequence[str], str]]
60class _ScaledGradientTensor:
61 """Wrapper around gradient tensors that applies backward_scale to sum operations.
63 This works around a PyTorch bug/behavior where multiplying gradient tensors
64 element-wise in backward hooks gives incorrect sums.
65 """
67 def __init__(self, tensor: Tensor, scale: float):
68 self._tensor = tensor
69 self._scale = scale
71 def sum(self, *args, **kwargs):
72 """Override sum to apply scaling to the result, not the tensor."""
73 result = self._tensor.sum(*args, **kwargs)
74 if isinstance(result, Tensor) and result.numel() == 1:
75 # Scalar result - apply scale
76 return result * self._scale
77 return result
79 def __getattr__(self, name):
80 """Delegate all other attributes to the wrapped tensor."""
81 return getattr(self._tensor, name)
83 def __repr__(self):
84 return f"ScaledGradientTensor({self._tensor}, scale={self._scale})"
87@runtime_checkable
88class _HookFunctionProtocol(Protocol):
89 """Protocol for hook functions."""
91 def __call__(self, tensor: Tensor, *, hook: "HookPoint") -> Union[Any, None]:
92 ...
95HookFunction = _HookFunctionProtocol # Callable[..., _HookFunctionProtocol]
97DeviceType = Optional[torch.device]
98_grad_t = Union[tuple[Tensor, ...], Tensor]
101class _AliasedHookPoint:
102 """
103 A lightweight wrapper that represents a HookPoint with an aliased name.
105 This is used when a hook is registered with multiple names (e.g., in compatibility mode
106 where both canonical and legacy names should trigger the hook). Instead of modifying
107 the original HookPoint's name, we create this wrapper that delegates to the original
108 HookPoint but presents a different name to the user's hook function.
109 """
111 def __init__(self, alias_name: str, target: "HookPoint"):
112 """
113 Create an aliased view of a HookPoint.
115 Args:
116 alias_name: The name to present to the hook function
117 target: The original HookPoint to delegate to
118 """
119 self._alias_name = alias_name
120 self._target = target
122 @property
123 def name(self) -> Optional[str]:
124 """Return the alias name."""
125 return self._alias_name
127 @property
128 def ctx(self) -> dict:
129 """Delegate to the target's context."""
130 return self._target.ctx
132 @property
133 def hook_conversion(self):
134 """Delegate to the target's hook conversion."""
135 return self._target.hook_conversion
137 def layer(self) -> int:
138 """
139 Extract layer index from the alias name.
141 Returns the layer index for hook names like 'blocks.0.attn.hook_pattern' -> 0
142 """
143 if self._alias_name is None:
144 raise ValueError("Name cannot be None")
145 split_name = self._alias_name.split(".")
146 return int(split_name[1])
149class HookPoint(nn.Module):
150 """
151 A helper class to access intermediate activations in a PyTorch model (inspired by Garcon).
153 HookPoint is a dummy module that acts as an identity function by default. By wrapping any
154 intermediate activation in a HookPoint, it provides a convenient way to add PyTorch hooks.
155 """
157 def __init__(self):
158 super().__init__()
159 self.fwd_hooks: list[LensHandle] = []
160 self.bwd_hooks: list[LensHandle] = []
161 self.ctx = {}
163 # A variable giving the hook's name (from the perspective of the root
164 # module) - this is set by the root module at setup.
165 self.name: Optional[str] = None
167 # Hook conversion for input and output transformations
168 self.hook_conversion: Optional[BaseTensorConversion] = None
170 # Backward gradient scale factor (for compatibility between architectures)
171 # This scales the SUM of gradients, not element-wise (to avoid PyTorch bugs)
172 self.backward_scale: float = 1.0
174 # Position axis of the activation in its batched layout, for ``pos_slice``.
175 # None means the default (dim 1, or -2 for attention patterns/scores); a
176 # component sets it when it fires a hook on a layout whose position axis
177 # is elsewhere, e.g. [batch, heads, pos, d_head] or [batch, channels, pos].
178 self.pos_dim: Optional[int] = None
180 def __repr__(self) -> str:
181 bits = [f"name={self.name!r}"] if self.name is not None else []
182 if self.fwd_hooks:
183 bits.append(f"{len(self.fwd_hooks)} fwd")
184 if self.bwd_hooks: 184 ↛ 185line 184 didn't jump to line 185 because the condition on line 184 was never true
185 bits.append(f"{len(self.bwd_hooks)} bwd")
186 return f"HookPoint({', '.join(bits)})" if bits else "HookPoint()"
188 def add_perma_hook(self, hook: HookFunction, dir: Literal["fwd", "bwd"] = "fwd") -> None:
189 self.add_hook(hook, dir=dir, is_permanent=True)
191 def add_hook(
192 self,
193 hook: HookFunction,
194 dir: Literal["fwd", "bwd"] = "fwd",
195 is_permanent: bool = False,
196 level: Optional[int] = None,
197 prepend: bool = False,
198 alias_names: Optional[list[str]] = None,
199 ) -> LensHandle:
200 """
201 Hook format is fn(activation, hook_name)
202 Change it into PyTorch hook format (this includes input and output,
203 which are the same for a HookPoint)
204 If prepend is True, add this hook before all other hooks
205 If alias_names is provided, the hook will be called once for each alias name,
206 receiving a temporary HookPoint-like object with that name instead of self
207 (useful for compatibility mode aliases)
208 Returns the LensHandle registered for this hook.
209 """
211 def full_hook(
212 module: torch.nn.Module,
213 module_input: Any,
214 module_output: Any,
215 ):
216 if (
217 dir == "bwd"
218 ): # For a backwards hook, module_output is a tuple of (grad,) - I don't know why.
219 module_output = module_output[0]
221 # Apply backward scaling if needed (wrap tensor to scale sum operations)
222 if self.backward_scale != 1.0: 222 ↛ 223line 222 didn't jump to line 223 because the condition on line 222 was never true
223 module_output = _ScaledGradientTensor(module_output, self.backward_scale)
225 # Apply input conversion if hook_conversion exists
226 if self.hook_conversion is not None:
227 module_output = self.hook_conversion.convert(module_output)
229 # Apply the hook for each name (or just once with canonical name)
230 if alias_names is not None:
231 # Call the hook once for each alias name
232 # Create a simple wrapper that acts like a HookPoint but with a different name
233 hook_result = None
234 hook_changed_output = False
235 for alias_name in alias_names:
236 # Create a view of this HookPoint with the alias name
237 hook_with_alias = _AliasedHookPoint(alias_name, self)
238 hook_result = hook(module_output, hook=hook_with_alias) # type: ignore[arg-type]
240 # If the hook modified the output, use that for subsequent calls
241 if hook_result is not None:
242 module_output = hook_result
243 hook_changed_output = True
244 if hook_changed_output: 244 ↛ 251line 244 didn't jump to line 251 because the condition on line 244 was always true
245 hook_result = module_output
246 else:
247 # Call the hook once with the canonical name (self)
248 hook_result = hook(module_output, hook=self)
250 # Apply output reversion if hook_conversion exists and hook returned a value
251 if hook_result is not None and self.hook_conversion is not None:
252 hook_result = self.hook_conversion.revert(hook_result)
254 # For backward hooks, PyTorch expects the return to be a tuple of (grad,)
255 if dir == "bwd" and hook_result is not None:
256 return (
257 hook_result
258 if isinstance(hook_result, tuple) and len(hook_result) == 1
259 else (hook_result,)
260 )
262 return hook_result
264 # annotate the `full_hook` with the string representation of the `hook` function
265 if isinstance(hook, partial):
266 # partial.__repr__() can be extremely slow if arguments contain large objects, which
267 # is common when caching tensors.
268 full_hook.__name__ = f"partial({hook.func.__repr__()},...)"
269 else:
270 full_hook.__name__ = hook.__repr__()
272 if dir == "fwd":
273 pt_handle = self.register_forward_hook(full_hook, prepend=prepend)
274 visible_hooks = self.fwd_hooks
275 elif dir == "bwd": 275 ↛ anywhereline 275 didn't jump anywhere: it always raised an exception.
276 # Wrap full_hook's bare Tensor return in tuple for PyTorch's backward API
277 def _bwd_hook_wrapper(
278 module: torch.nn.Module,
279 grad_input: Any,
280 grad_output: Any,
281 ):
282 result = full_hook(module, grad_input, grad_output)
283 if result is None:
284 return None
285 if isinstance(result, tuple): 285 ↛ 287line 285 didn't jump to line 287 because the condition on line 285 was always true
286 return result
287 return (result,)
289 if isinstance(hook, partial):
290 _bwd_hook_wrapper.__name__ = f"partial({hook.func.__repr__()},...)"
291 else:
292 _bwd_hook_wrapper.__name__ = hook.__repr__()
293 pt_handle = self.register_full_backward_hook(_bwd_hook_wrapper, prepend=prepend)
294 visible_hooks = self.bwd_hooks
295 else:
296 raise ValueError(f"Invalid direction {dir}")
298 handle = LensHandle(pt_handle, is_permanent, level, user_hook=hook)
300 if prepend:
301 # we could just pass this as an argument in PyTorch 2.0, but for now we manually do this...
302 visible_hooks.insert(0, handle)
304 else:
305 visible_hooks.append(handle)
306 return handle
308 def has_hooks(
309 self,
310 dir: Literal["fwd", "bwd", "both"] = "both",
311 including_permanent: bool = True,
312 level: Optional[int] = None,
313 ) -> bool:
314 """Check if this HookPoint has any active hooks.
316 Args:
317 dir: Direction of hooks to check ("fwd", "bwd", or "both")
318 including_permanent: Whether to include permanent hooks in the check
319 level: Only check hooks at this context level (None for all levels)
321 Returns:
322 True if any matching hooks are found, False otherwise
323 """
325 def _has_hooks_in_direction(handles: list[LensHandle]) -> bool:
326 return any(_handle_matches(h, including_permanent, level) for h in handles)
328 if dir == "fwd":
329 return _has_hooks_in_direction(self.fwd_hooks)
330 elif dir == "bwd":
331 return _has_hooks_in_direction(self.bwd_hooks)
332 elif dir == "both": 332 ↛ 337line 332 didn't jump to line 337 because the condition on line 332 was always true
333 return _has_hooks_in_direction(self.fwd_hooks) or _has_hooks_in_direction(
334 self.bwd_hooks
335 )
336 else:
337 raise ValueError(f"Invalid direction {dir}")
339 def remove_hooks(
340 self,
341 dir: Literal["fwd", "bwd", "both"] = "fwd",
342 including_permanent: bool = False,
343 level: Optional[int] = None,
344 ) -> None:
345 def _remove_hooks(handles: list[LensHandle]) -> list[LensHandle]:
346 output_handles = []
347 for handle in handles:
348 if not _handle_matches(handle, including_permanent, level):
349 output_handles.append(handle)
350 continue
351 handle.hook.remove()
352 return output_handles
354 if dir == "fwd" or dir == "both":
355 self.fwd_hooks = _remove_hooks(self.fwd_hooks)
356 if dir == "bwd" or dir == "both":
357 self.bwd_hooks = _remove_hooks(self.bwd_hooks)
358 if dir not in ["fwd", "bwd", "both"]: 358 ↛ 359line 358 didn't jump to line 359 because the condition on line 358 was never true
359 raise ValueError(f"Invalid direction {dir}")
361 def clear_context(self):
362 del self.ctx
363 self.ctx = {}
365 def enable_reshape(
366 self,
367 hook_conversion: Optional[BaseTensorConversion] = None,
368 ) -> None:
369 """
370 Enable reshape functionality for this hook point using a BaseTensorConversion.
372 Args:
373 hook_conversion: BaseTensorConversion instance to handle input/output transformations.
374 The convert() method will be used for input transformation,
375 and the revert() method will be used for output transformation.
376 """
377 self.hook_conversion = hook_conversion
379 def forward(self, x: Tensor) -> Tensor:
380 return x
382 def layer(self):
383 # Returns the layer index if the name has the form 'blocks.{layer}.{...}'
384 # Helper for models whose hook names follow the blocks.{layer}.* scheme
385 # If it doesn't have this form, raises an error -
386 if self.name is None:
387 raise ValueError("Name cannot be None")
388 split_name = self.name.split(".")
389 return int(split_name[1])
392# %%
393class HookIntrospectionMixin:
394 """``list_hooks()`` mixins for any class exposing a ``hook_dict``.
396 Accessed via ``getattr`` so subclasses can provide ``hook_dict`` as either
397 an instance attribute (``HookedRootModule``) or a ``@property`` (``TransformerBridge``).
398 """
400 def list_hooks(
401 self,
402 name_filter: NamesFilter = None,
403 dir: Literal["fwd", "bwd", "both"] = "both",
404 including_permanent: bool = True,
405 ) -> dict[str, list[LensHandle]]:
406 """Return attached hooks grouped by HookPoint name; empty HookPoints are omitted.
408 Args:
409 name_filter: A hook name, list of names, or predicate. ``None`` matches all.
410 dir: Restrict to forward, backward, or both directions.
411 including_permanent: If False, drop permanent hooks from the result.
412 """
413 if name_filter is None:
414 matches: Callable[[str], bool] = lambda _: True
415 elif callable(name_filter):
416 matches = name_filter
417 elif isinstance(name_filter, str):
418 target = name_filter
419 matches = lambda n: n == target
420 else:
421 allowed = set(name_filter)
422 matches = lambda n: n in allowed
424 out: dict[str, list[LensHandle]] = {}
425 hook_dict: dict[str, HookPoint] = getattr(self, "hook_dict")
426 for name, hp in hook_dict.items():
427 if not matches(name):
428 continue
429 handles: list[LensHandle] = []
430 if dir in ("fwd", "both"):
431 handles.extend(hp.fwd_hooks)
432 if dir in ("bwd", "both"):
433 handles.extend(hp.bwd_hooks)
434 handles = [h for h in handles if _handle_matches(h, including_permanent, None)]
435 if handles:
436 out[name] = handles
437 return out
440# HookedRootModule moved to transformer_lens.HookedRootModule in 3.0; the
441# deprecated re-export from here was removed in 4.0 as promised. The class itself
442# is kept — import it from transformer_lens or transformer_lens.HookedRootModule.
443def __getattr__(name: str):
444 if name == "HookedRootModule":
445 # ImportError, not AttributeError: the from-import machinery discards an
446 # AttributeError and raises its own unchained "cannot import name", so
447 # `from transformer_lens.hook_points import HookedRootModule` would lose this
448 # pointer. ImportError reaches both attribute access and from-import.
449 raise ImportError(
450 "Importing HookedRootModule from transformer_lens.hook_points was "
451 "removed in TransformerLens 4.0. Import it from transformer_lens "
452 "(preferred) or transformer_lens.HookedRootModule instead — the class "
453 "itself is kept.",
454 name=name,
455 )
456 raise AttributeError(f"module {__name__!r} has no attribute {name!r}")