Skip to content

interlens.interp.capture

interlens.interp.capture

CaptureRequest dataclass 

CaptureRequest(cache: ActivationCache, spec: CaptureSpec)

A pending capture handed to generate: where to store records (cache) and what to grab (spec).

The Conversation builds this (via conv.capture(...)) and the participant fills the cache with records tagged by participant + turn, so the caller ends up with structurally-tagged activations.

CapturedSite

Bases: NamedTuple

One activation captured by capture_activations: the tensor ([seq, d_model]) at a given layer and site. A lightweight typed row (unpacks like the old (layer, site, tensor) tuple) that the participant folds into a fully-tagged ActivationRecord.

capture_activations 

capture_activations(
    model: "PreTrainedModel",
    input_ids: Tensor,
    spec: CaptureSpec,
) -> list[CapturedSite]

Run one clean forward pass over input_ids and return [(layer, site, tensor[seq, d_model])].

Design choice: capture is a separate forward pass over the full (prompt + generated) sequence rather than accumulating hooks across the multi-step decode loop. This is simpler and provably complete — every position is present in one pass — at the cost of one extra forward. Residual-stream activations come from output_hidden_states (no hooks needed); attn/mlp sublayer outputs come from forward hooks on the corresponding submodules.

Note: attn captures the attention sublayer output (post-o_proj), which is available under any attention backend. Attention weights/patterns are NOT captured here — those require attn_implementation='eager' + output_attentions and are out of scope for the default kernel.

Source code in src/interlens/interp/capture.py 
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def capture_activations(model: "PreTrainedModel", input_ids: torch.Tensor, spec: CaptureSpec) -> list[CapturedSite]:
	"""Run one clean forward pass over ``input_ids`` and return ``[(layer, site, tensor[seq, d_model])]``.

	Design choice: capture is a *separate forward pass* over the full (prompt + generated) sequence rather than
	accumulating hooks across the multi-step decode loop. This is simpler and provably complete — every position
	is present in one pass — at the cost of one extra forward. Residual-stream activations come from
	``output_hidden_states`` (no hooks needed); ``attn``/``mlp`` sublayer outputs come from forward hooks on the
	corresponding submodules.

	Note: ``attn`` captures the attention *sublayer output* (post-o_proj), which is available under any attention
	backend. Attention *weights/patterns* are NOT captured here — those require ``attn_implementation='eager'`` +
	``output_attentions`` and are out of scope for the default kernel.
	"""
	layers = decoder_layers(model)
	want = spec.layers if spec.layers is not None else tuple(range(len(layers)))
	sites = set(spec.sites)

	results: list[CapturedSite] = []
	handles = []
	hook_store: dict[tuple[int, Site], torch.Tensor] = {}

	def make_hook(layer_idx: int, site: Site):
		def hook(module, inputs, output):
			hs = output[0] if isinstance(output, tuple) else output
			hook_store[(layer_idx, site)] = hs.detach()[0]  # [seq, d_model]
		return hook

	for li in want:
		if "attn" in sites:
			handles.append(layers[li].self_attn.register_forward_hook(make_hook(li, "attn")))
		if "mlp" in sites:
			handles.append(layers[li].mlp.register_forward_hook(make_hook(li, "mlp")))

	need_hidden = "residual" in sites
	try:
		with torch.inference_mode():
			out = model(input_ids, output_hidden_states=need_hidden, use_cache=False)
	finally:
		for h in handles:
			h.remove()

	if need_hidden:
		# hidden_states is (embeddings, layer_0_out, ..., layer_{L-1}_out); layer li output is index li+1.
		hs = out.hidden_states
		for li in want:
			results.append(CapturedSite(li, "residual", hs[li + 1][0]))
	for (li, site), tensor in hook_store.items():
		results.append(CapturedSite(li, site, tensor))

	return results