Coverage for transformer_lens/model_bridge/sources/inspect/hooks.py: 100%

1"""Canonical hook names ↔ (layer, kind) for the Inspect HF provider. 

2 

3Torch-free; shared by the provider (capture/intervene) and the driver (supported set, 

4decode). Covers the ``d_model``-shaped decoder-layer boundaries plus the head-split 

5attention hooks (q/k/v/z, pattern) where the provider's structural probe finds the 

6projections; ``attn_scores``, ``embed``, and ``ln_final`` (fold-LN convention) stay 

7non-fireable. 

8 

9Names are TransformerBridge-native (``blocks.{i}.hook_out``, ``.attn.hook_out``, ...), 

10not the HookedTransformer aliases. A bridge cache carries both with identical values, 

11so parity vs ``boot_transformers`` still resolves. 

12 

13Which boundaries are actually fireable is decided per-model by the provider's structural 

14self-check, not a hand-kept architecture list: it locates attn/mlp and probes whether the 

15``resid_pre + attn_out`` derivation holds, gating ``resid_mid`` otherwise; head-split 

16kinds need separate q/k/v projections (q/k/v), a locatable out-projection (z), or eager 

17attention (pattern). ``supported_hook_points(n_layers, kinds=...)`` filters to that set. 

18""" 

19from __future__ import annotations 

20 

21import re 

22from typing import Iterable, Optional 

23 

24ALL_KINDS = frozenset({"resid_pre", "resid_mid", "resid_post", "attn_out", "mlp_out"}) 

25# Head-split attention kinds, served only when structurally detected (never by default): 

26# q/k/v are the pre-RoPE projection outputs, z is the out-projection input (all 

27# ``(seq, heads, d_head)``); pattern is post-softmax attention ``(heads, q_pos, k_pos)``. 

28HEAD_KINDS = frozenset({"q", "k", "v", "z", "pattern"}) 

29 

30# One canonical TransformerBridge name per boundary (no aliases — avoids duplicate 

31# HookPoints/cache entries). resid_mid (ln2.hook_in) is derived (resid_pre + 

32# attn_out), so it's capture-only. 

33_KIND_NAMES = { 

34 "resid_pre": "blocks.{i}.hook_in", 

35 "resid_mid": "blocks.{i}.ln2.hook_in", 

36 "resid_post": "blocks.{i}.hook_out", 

37 "attn_out": "blocks.{i}.attn.hook_out", 

38 "mlp_out": "blocks.{i}.mlp.hook_out", 

39 "q": "blocks.{i}.attn.hook_q", 

40 "k": "blocks.{i}.attn.hook_k", 

41 "v": "blocks.{i}.attn.hook_v", 

42 "z": "blocks.{i}.attn.hook_z", 

43 "pattern": "blocks.{i}.attn.hook_pattern", 

44} 

45# pattern is read from the forward's output_attentions — nothing to write back, so it's 

46# capture-only (like the derived resid_mid). 

47INTERVENEABLE_KINDS = frozenset( 

48 {"resid_pre", "attn_out", "mlp_out", "resid_post", "q", "k", "v", "z"} 

49) 

50 

51# Rank of each kind's batchless wire array; the driver unsqueezes exactly one batch dim. 

52WIRE_BATCHLESS_NDIM = { 

53 **{kind: 2 for kind in ALL_KINDS}, # (seq, d_model) 

54 "q": 3, 

55 "k": 3, 

56 "v": 3, 

57 "z": 3, # (seq, heads, d_head) 

58 "pattern": 3, # (heads, q_pos, k_pos) 

59} 

60 

61_SUFFIX_TO_KIND = { 

62 "hook_in": "resid_pre", 

63 "ln2.hook_in": "resid_mid", 

64 "hook_out": "resid_post", 

65 "attn.hook_out": "attn_out", 

66 "mlp.hook_out": "mlp_out", 

67 "attn.hook_q": "q", 

68 "attn.hook_k": "k", 

69 "attn.hook_v": "v", 

70 "attn.hook_z": "z", 

71 "attn.hook_pattern": "pattern", 

72} 

73_BLOCK = re.compile(r"^blocks\.(\d+)\.(.+)$") 

74 

75 

76def supported_hook_points(n_layers: int, kinds: Optional[Iterable[str]] = None) -> frozenset[str]: 

77 """Fireable hook names across all layers. ``kinds=None`` means all *boundary* kinds 

78 (head-split kinds are opt-in — a provider must detect and list them explicitly); 

79 pass the provider's detected kinds to gate (e.g. drop ``resid_mid`` for parallel).""" 

80 selected = ALL_KINDS if kinds is None else kinds 

81 return frozenset(_KIND_NAMES[k].format(i=i) for i in range(n_layers) for k in selected) 

82 

83 

84def all_hook_points(n_layers: int) -> frozenset[str]: 

85 """Every hook name the registry can serve (boundaries + head-split) — the universe a 

86 driver subtracts its supported set from to build ``non_fireable_hook_points``.""" 

87 return supported_hook_points(n_layers, ALL_KINDS | HEAD_KINDS) 

88 

89 

90def nonfireable_hook_points(n_layers: int) -> frozenset[str]: 

91 """Hooks no provider configuration can fire (embed, ln_final, pre-softmax scores). 

92 Head-split q/k/v/z/pattern are *conditionally* fireable and belong here only when a 

93 model's structural probe gates them — the driver handles that subtraction.""" 

94 names = ["embed.hook_out", "ln_final.hook_normalized", "unembed.hook_out"] 

95 names += [f"blocks.{i}.attn.hook_attn_scores" for i in range(n_layers)] 

96 return frozenset(names) 

97 

98 

99def resolve(name: str) -> tuple[int, str] | None: 

100 """Canonical hook name → ``(layer, kind)``, or ``None`` if not a fireable hook.""" 

101 match = _BLOCK.match(name) 

102 if match is None: 

103 return None 

104 kind = _SUFFIX_TO_KIND.get(match.group(2)) 

105 return (int(match.group(1)), kind) if kind is not None else None 

106 

107 

108def wire_key(layer: int, kind: str) -> str: 

109 """Stable key for one captured boundary in the activation payload.""" 

110 return f"{layer}:{kind}" 

111 

112 

113def name_from_wire_key(key: str) -> str | None: 

114 """Inverse of ``wire_key`` ∘ ``resolve``: ``"<layer>:<kind>"`` → the canonical hook 

115 name, or ``None`` if the kind is unknown.""" 

116 layer, _, kind = key.partition(":") 

117 template = _KIND_NAMES.get(kind) 

118 return template.format(i=int(layer)) if template and layer.isdigit() else None