Coverage for transformer_lens/hook_points.py: 86%
200 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +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 # Position axis of the activation in its batched layout, for ``pos_slice``.
168 # None means the default (dim 1, or -2 for attention patterns/scores); a
169 # component sets it when it fires a hook on a layout whose position axis
170 # is elsewhere, e.g. [batch, heads, pos, d_head] or [batch, channels, pos].
171 self.pos_dim: Optional[int] = None
173 def __repr__(self) -> str:
174 bits = [f"name={self.name!r}"] if self.name is not None else []
175 if self.fwd_hooks:
176 bits.append(f"{len(self.fwd_hooks)} fwd")
177 if self.bwd_hooks: 177 ↛ 178line 177 didn't jump to line 178 because the condition on line 177 was never true
178 bits.append(f"{len(self.bwd_hooks)} bwd")
179 return f"HookPoint({', '.join(bits)})" if bits else "HookPoint()"
181 def add_perma_hook(self, hook: HookFunction, dir: Literal["fwd", "bwd"] = "fwd") -> None:
182 self.add_hook(hook, dir=dir, is_permanent=True)
184 def add_hook(
185 self,
186 hook: HookFunction,
187 dir: Literal["fwd", "bwd"] = "fwd",
188 is_permanent: bool = False,
189 level: Optional[int] = None,
190 prepend: bool = False,
191 alias_names: Optional[list[str]] = None,
192 ) -> None:
193 """
194 Hook format is fn(activation, hook_name)
195 Change it into PyTorch hook format (this includes input and output,
196 which are the same for a HookPoint)
197 If prepend is True, add this hook before all other hooks
198 If alias_names is provided, the hook will be called once for each alias name,
199 receiving a temporary HookPoint-like object with that name instead of self
200 (useful for compatibility mode aliases)
201 """
203 def full_hook(
204 module: torch.nn.Module,
205 module_input: Any,
206 module_output: Any,
207 ):
208 if (
209 dir == "bwd"
210 ): # For a backwards hook, module_output is a tuple of (grad,) - I don't know why.
211 module_output = module_output[0]
213 # Apply backward scaling if needed (wrap tensor to scale sum operations)
214 if self.backward_scale != 1.0: 214 ↛ 215line 214 didn't jump to line 215 because the condition on line 214 was never true
215 module_output = _ScaledGradientTensor(module_output, self.backward_scale)
217 # Apply input conversion if hook_conversion exists
218 if self.hook_conversion is not None:
219 module_output = self.hook_conversion.convert(module_output)
221 # Apply the hook for each name (or just once with canonical name)
222 if alias_names is not None:
223 # Call the hook once for each alias name
224 # Create a simple wrapper that acts like a HookPoint but with a different name
225 hook_result = None
226 hook_changed_output = False
227 for alias_name in alias_names:
228 # Create a view of this HookPoint with the alias name
229 hook_with_alias = _AliasedHookPoint(alias_name, self)
230 hook_result = hook(module_output, hook=hook_with_alias) # type: ignore[arg-type]
232 # If the hook modified the output, use that for subsequent calls
233 if hook_result is not None:
234 module_output = hook_result
235 hook_changed_output = True
236 if hook_changed_output: 236 ↛ 243line 236 didn't jump to line 243 because the condition on line 236 was always true
237 hook_result = module_output
238 else:
239 # Call the hook once with the canonical name (self)
240 hook_result = hook(module_output, hook=self)
242 # Apply output reversion if hook_conversion exists and hook returned a value
243 if hook_result is not None and self.hook_conversion is not None:
244 hook_result = self.hook_conversion.revert(hook_result)
246 # For backward hooks, PyTorch expects the return to be a tuple of (grad,)
247 if dir == "bwd" and hook_result is not None:
248 return (
249 hook_result
250 if isinstance(hook_result, tuple) and len(hook_result) == 1
251 else (hook_result,)
252 )
254 return hook_result
256 # annotate the `full_hook` with the string representation of the `hook` function
257 if isinstance(hook, partial):
258 # partial.__repr__() can be extremely slow if arguments contain large objects, which
259 # is common when caching tensors.
260 full_hook.__name__ = f"partial({hook.func.__repr__()},...)"
261 else:
262 full_hook.__name__ = hook.__repr__()
264 if dir == "fwd":
265 pt_handle = self.register_forward_hook(full_hook, prepend=prepend)
266 visible_hooks = self.fwd_hooks
267 elif dir == "bwd": 267 ↛ anywhereline 267 didn't jump anywhere: it always raised an exception.
268 # Wrap full_hook's bare Tensor return in tuple for PyTorch's backward API
269 def _bwd_hook_wrapper(
270 module: torch.nn.Module,
271 grad_input: Any,
272 grad_output: Any,
273 ):
274 result = full_hook(module, grad_input, grad_output)
275 if result is None:
276 return None
277 if isinstance(result, tuple): 277 ↛ 279line 277 didn't jump to line 279 because the condition on line 277 was always true
278 return result
279 return (result,)
281 if isinstance(hook, partial):
282 _bwd_hook_wrapper.__name__ = f"partial({hook.func.__repr__()},...)"
283 else:
284 _bwd_hook_wrapper.__name__ = hook.__repr__()
285 pt_handle = self.register_full_backward_hook(_bwd_hook_wrapper, prepend=prepend)
286 visible_hooks = self.bwd_hooks
287 else:
288 raise ValueError(f"Invalid direction {dir}")
290 handle = LensHandle(pt_handle, is_permanent, level, user_hook=hook)
292 if prepend:
293 # we could just pass this as an argument in PyTorch 2.0, but for now we manually do this...
294 visible_hooks.insert(0, handle)
296 else:
297 visible_hooks.append(handle)
299 def has_hooks(
300 self,
301 dir: Literal["fwd", "bwd", "both"] = "both",
302 including_permanent: bool = True,
303 level: Optional[int] = None,
304 ) -> bool:
305 """Check if this HookPoint has any active hooks.
307 Args:
308 dir: Direction of hooks to check ("fwd", "bwd", or "both")
309 including_permanent: Whether to include permanent hooks in the check
310 level: Only check hooks at this context level (None for all levels)
312 Returns:
313 True if any matching hooks are found, False otherwise
314 """
316 def _has_hooks_in_direction(handles: list[LensHandle]) -> bool:
317 for handle in handles:
318 # Check if this hook matches our criteria
319 if not including_permanent and handle.is_permanent:
320 continue
321 if level is not None and handle.context_level != level:
322 continue
323 return True
324 return False
326 if dir == "fwd":
327 return _has_hooks_in_direction(self.fwd_hooks)
328 elif dir == "bwd":
329 return _has_hooks_in_direction(self.bwd_hooks)
330 elif dir == "both": 330 ↛ 335line 330 didn't jump to line 335 because the condition on line 330 was always true
331 return _has_hooks_in_direction(self.fwd_hooks) or _has_hooks_in_direction(
332 self.bwd_hooks
333 )
334 else:
335 raise ValueError(f"Invalid direction {dir}")
337 def remove_hooks(
338 self,
339 dir: Literal["fwd", "bwd", "both"] = "fwd",
340 including_permanent: bool = False,
341 level: Optional[int] = None,
342 ) -> None:
343 def _remove_hooks(handles: list[LensHandle]) -> list[LensHandle]:
344 output_handles = []
345 for handle in handles:
346 if including_permanent:
347 handle.hook.remove()
348 elif (not handle.is_permanent) and (level is None or handle.context_level == level):
349 handle.hook.remove()
350 else:
351 output_handles.append(handle)
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 if not including_permanent:
435 handles = [h for h in handles if not h.is_permanent]
436 if handles:
437 out[name] = handles
438 return out
441# HookedRootModule moved to transformer_lens.HookedRootModule in 3.0; the
442# deprecated re-export from here was removed in 4.0 as promised. The class itself
443# is kept — import it from transformer_lens or transformer_lens.HookedRootModule.
444def __getattr__(name: str):
445 if name == "HookedRootModule":
446 # ImportError, not AttributeError: the from-import machinery discards an
447 # AttributeError and raises its own unchained "cannot import name", so
448 # `from transformer_lens.hook_points import HookedRootModule` would lose this
449 # pointer. ImportError reaches both attribute access and from-import.
450 raise ImportError(
451 "Importing HookedRootModule from transformer_lens.hook_points was "
452 "removed in TransformerLens 4.0. Import it from transformer_lens "
453 "(preferred) or transformer_lens.HookedRootModule instead — the class "
454 "itself is kept.",
455 name=name,
456 )
457 raise AttributeError(f"module {__name__!r} has no attribute {name!r}")