spb/wpai0001_uqo

1"""Hook points and interventions for activation capture, patching and steering.23Design contract (tested in tests/causal_sanity/test_hook_noop.py):45* A ``HookManager`` with ``enabled=False`` or with no registered interventions returns the input6  tensor object itself (identity, no copy), so a model instrumented with hook points produces7  bit-identical outputs to the uninstrumented computation.8* Interventions are pure functions ``(tensor, ctx) -> tensor`` applied in registration order at a9  named hook point. ``ctx`` carries the hook name and the batch of token positions.10* Capture interventions record (detached, cloned) activations without modifying the stream.1112Hook names used by ``ToyTransformer``:13  ``embed``, ``resid_pre.{l}``, ``attn_heads.{l}`` ([B,T,H,dh] before the output projection),14  ``attn_out.{l}``, ``resid_mid.{l}``, ``mlp_out.{l}``, ``resid_post.{l}``, ``final_norm``.15"""16from __future__ import annotations1718from dataclasses import dataclass, field19from typing import Callable, Protocol2021import torch2223Tensor = torch.Tensor242526@dataclass27class HookContext:28    name: str29    positions: Tensor | None = None  # [B] answer positions when available30    extra: dict = field(default_factory=dict)313233class Intervention(Protocol):34    def __call__(self, x: Tensor, ctx: HookContext) -> Tensor: ...353637class HookManager:38    def __init__(self) -> None:39        self.enabled: bool = True40        self._hooks: dict[str, list[Intervention]] = {}41        self.calls: int = 04243    # -- registration -------------------------------------------------------------------44    def add(self, name: str, fn: Intervention) -> None:45        self._hooks.setdefault(name, []).append(fn)4647    def remove(self, name: str, fn: Intervention | None = None) -> None:48        if fn is None:49            self._hooks.pop(name, None)50        elif name in self._hooks:51            self._hooks[name] = [f for f in self._hooks[name] if f is not fn]5253    def clear(self) -> None:54        self._hooks.clear()5556    def names(self) -> list[str]:57        return [n for n, fns in self._hooks.items() if fns]5859    # -- application --------------------------------------------------------------------60    def __call__(self, name: str, x: Tensor, ctx: HookContext | None = None) -> Tensor:61        if not self.enabled:62            return x63        fns = self._hooks.get(name)64        if not fns:65            return x66        self.calls += 167        c = ctx or HookContext(name=name)68        c.name = name69        for fn in fns:70            x = fn(x, c)71        return x7273    class _Disabled:74        def __init__(self, mgr: "HookManager") -> None:75            self.mgr, self.prev = mgr, mgr.enabled7677        def __enter__(self) -> None:78            self.mgr.enabled = False7980        def __exit__(self, *a) -> None:81            self.mgr.enabled = self.prev8283    def disabled(self) -> "HookManager._Disabled":84        return HookManager._Disabled(self)858687# ---------------------------------------------------------------------------------------88# Built-in interventions89# ---------------------------------------------------------------------------------------90class Capture:91    """Record activations (detached clone) without modifying the stream."""9293    def __init__(self) -> None:94        self.store: dict[str, Tensor] = {}9596    def __call__(self, x: Tensor, ctx: HookContext) -> Tensor:97        self.store[ctx.name] = x.detach().clone()98        return x99100101def _pos_mask(x: Tensor, positions: Tensor | list[int] | slice | None) -> Tensor:102    """Boolean mask [B,T] selecting token positions (per-example positions or shared slice)."""103    B, T = x.shape[0], x.shape[1]104    mask = torch.zeros(B, T, dtype=torch.bool, device=x.device)105    if positions is None:106        mask[:] = True107    elif isinstance(positions, slice):108        mask[:, positions] = True109    elif isinstance(positions, torch.Tensor) and positions.ndim == 1 and positions.shape[0] == B:110        mask[torch.arange(B, device=x.device), positions.to(x.device)] = True111    else:112        mask[:, list(positions)] = True  # type: ignore[arg-type]113    return mask114115116class Replace:117    """Activation patching: overwrite selected positions with ``source`` values of the same shape."""118119    def __init__(self, source: Tensor, positions=None, heads: list[int] | None = None, dims: slice | None = None) -> None:120        self.source, self.positions, self.heads, self.dims = source, positions, heads, dims121122    def __call__(self, x: Tensor, ctx: HookContext) -> Tensor:123        src = self.source.to(x.device, x.dtype)124        mask = _pos_mask(x, self.positions)125        out = x.clone()126        if self.heads is not None:  # x: [B,T,H,dh]127            m = mask[:, :, None, None].expand_as(x).clone()128            keep = torch.zeros(x.shape[2], dtype=torch.bool, device=x.device)129            keep[self.heads] = True130            m &= keep[None, None, :, None]131            out = torch.where(m, src, x)132        elif self.dims is not None:133            m = mask[:, :, None].expand_as(x).clone()134            dmask = torch.zeros(x.shape[-1], dtype=torch.bool, device=x.device)135            dmask[self.dims] = True136            m &= dmask[None, None, :]137            out = torch.where(m, src, x)138        else:139            m = mask.reshape(mask.shape + (1,) * (x.ndim - 2)).expand_as(x)140            out = torch.where(m, src, x)141        return out142143144class Add:145    """Additive steering: x[positions] += scale * vector (vector broadcast over the last dim)."""146147    def __init__(self, vector: Tensor, scale: float = 1.0, positions=None) -> None:148        self.vector, self.scale, self.positions = vector, scale, positions149150    def __call__(self, x: Tensor, ctx: HookContext) -> Tensor:151        v = self.vector.to(x.device, x.dtype)152        mask = _pos_mask(x, self.positions)153        m = mask.reshape(mask.shape + (1,) * (x.ndim - 2)).to(x.dtype)154        return x + self.scale * m * v155156157class Zero:158    def __init__(self, positions=None, heads: list[int] | None = None) -> None:159        self.positions, self.heads = positions, heads160161    def __call__(self, x: Tensor, ctx: HookContext) -> Tensor:162        return Replace(torch.zeros_like(x), self.positions, heads=self.heads)(x, ctx)163164165class GaussianNoise:166    """Causal-tracing style corruption with a fixed generator for reproducibility."""167168    def __init__(self, std: float, positions=None, seed: int = 0) -> None:169        self.std, self.positions, self.seed = std, positions, seed170171    def __call__(self, x: Tensor, ctx: HookContext) -> Tensor:172        g = torch.Generator(device="cpu").manual_seed(self.seed)173        noise = torch.randn(x.shape, generator=g, dtype=torch.float32).to(x.device, x.dtype) * self.std174        mask = _pos_mask(x, self.positions)175        m = mask.reshape(mask.shape + (1,) * (x.ndim - 2)).to(x.dtype)176        return x + m * noise177178179class ProjectSubspace:180    """Remove (mode='ablate') or keep only (mode='keep') the component in span(basis)."""181182    def __init__(self, basis: Tensor, mode: str = "ablate", positions=None) -> None:183        assert mode in {"ablate", "keep"}184        q, _ = torch.linalg.qr(basis.T.float() if basis.shape[0] < basis.shape[1] else basis.float())185        self.q, self.mode, self.positions = q, mode, positions  # [d, r] orthonormal columns186187    def __call__(self, x: Tensor, ctx: HookContext) -> Tensor:188        q = self.q.to(x.device, x.dtype)189        proj = (x @ q) @ q.T190        new = x - proj if self.mode == "ablate" else proj191        mask = _pos_mask(x, self.positions)192        m = mask.reshape(mask.shape + (1,) * (x.ndim - 2)).expand_as(x)193        return torch.where(m, new, x)194195196class Sequence:197    def __init__(self, *fns: Callable[[Tensor, HookContext], Tensor]) -> None:198        self.fns = fns199200    def __call__(self, x: Tensor, ctx: HookContext) -> Tensor:201        for f in self.fns:202            x = f(x, ctx)203        return x204