Intervention API Reference¶

This page documents the public intervention API: selectors, edit helpers, hooks, replay and bundles.

To apply one intervention to many forwards or a generate() loop, use spec.bind(model) (capture-free, about the cost of a plain hook) or tl.record(..., intervene=spec) (the model output is recording.output) when you need activations too, or rerun a steered trace with trace.run(model, x) (the guarded fast engine on a selective-save trace, with the spec applied when it targets modules; see the glossary entry "Guarded fast rerun"). x, like the inputs of tl.record, takes every input form tl.trace takes: a tensor, a positional list or tuple, and for a Hugging Face model a string, PIL image or modality dict ({"text": ..., "image": ...}), resolved with the model's own tokenizer or AutoProcessor exactly as the trace was. Which one to pick for a loop or for generation: Performance: choosing a fast path.

Selectors¶

Selectors resolve against completed Trace.layers records.

One interpreter evaluates every selector in every lifecycle (capture-time save=, post-hoc find_sites, live hooks); label_contains is case-insensitive and label_matches case-sensitive everywhere. Post-hoc label_contains/label_matches search the final layer_label only; exact tl.label additionally matches raw, short, and pass-qualified spellings. Capture-only selectors (tl.followed_by, tl.preceded_by) and mutator-only selectors (tl.facet, tl.head) refuse unsupported lifecycles with the typed SelectorCapabilityError (a SiteResolutionError subclass) instead of a generic message — upfront, before any per-site evaluation, for post-hoc resolution and live hook attachment alike.

Composites (&, |, ~) SHORT-CIRCUIT per site in every lifecycle: a tl.where predicate is only invoked for sites its siblings have not already decided, so predicates must not rely on side effects from seeing every site. &/| build nested binary composites, and deserialized target specs may carry flat n-ary child tuples; both shapes evaluate identically. Conjunctions are association-insensitive for the temporal sugar — a & tl.followed_by(x) & b behaves exactly like a & b & tl.followed_by(x) and the flat three-child spec — and degenerate arities keep identity semantics in evaluation and spec round-trips alike: an empty and matches everything, an empty or matches nothing, and a unary composite matches like its child. tl.followed_by/tl.preceded_by target specs serialize a selector inner structurally at every save level; an opaque-callable inner remains audit-only and refuses typed at rebuild rather than being reconstructed lossily.

Selectors compose with & and | for in-memory discovery:

candidate = tl.in_module("block") & tl.func("relu")
sites = log.find_sites(candidate, max_fanout=8)
exact = tl.label(sites.labels()[0])

The same selector objects can drive capture-time save and intervention predicates:

saved = tl.trace(model, x, save=tl.func("relu"))
windowed = tl.trace(
    model,
    x,
    save=tl.func("conv2d") & tl.followed_by(tl.func("relu")),
capture=CaptureOptions(lookback=4, lookback_payload_policy="detached_raw"),
)
patched = tl.trace(
    model,
    x,
    save=tl.func("relu"),
    intervene=tl.when(tl.func("relu"), tl.scale(0.5)),
)
conditional = tl.trace(
    model,
    x,
    intervene=tl.when(tl.func("relu") & tl.preceded_by(tl.in_module("block")), tl.zero_ablate()),
capture=CaptureOptions(lookback=4),
)

For saved or replayed specs, prefer exact tl.label(...) or another simple portable selector after discovery.

Forward Helpers¶

Forward helpers return Edit objects that can be passed to do (on the replay, rerun or set-only engine), or live capture hooks=....

The packaged SAE splice experiment, torchlens.bridge.sae.splice(...) (experimental), composes splice_module with fork + do: it swaps a duck-typed SAE's reconstruction (encode/decode pair; no SAE package required) back into the model at one site, replays downstream, and reports reconstruction fidelity plus output-level causal effect. The optional latents_edit= callable transforms the encoded latents before decoding -- the feature-level causal-testing knob. See notebooks/sae_splice_tutorial.ipynb for the end-to-end experiment.

Gradient Sites¶

A gradient is addressed as a site, not chosen by a keyword. tl.grad("relu_3") names the gradient of relu_3's OUTPUT (the tensor trace["relu_3"].grad stores), a bare "relu_3" names its activation, and tl.grad("relu_3") | "relu_3" names both. tl.grad is an address; it never calls autograd. It wraps a forward site once and refuses node selectors (tl.grad_fn, ...), facets and selections.

There are no gradient-only helpers. The ordinary edits apply to whichever tensor the site names:

Edits that read recorded forward state (tl.patch_from, tl.mean_ablate, tl.splice_module, torchlens.intervention._helpers.swap_with with a site or trace donor) refuse on a gradient site with edit_signal_mismatch.

trace.do(tl.grad(site), edit, model=model, inputs=x) reruns the forward with the edit armed; the next result.backward(loss) applies it. loss=lambda t: ... (a callable of the returned trace) does both in one call. The audit row status is "armed" until a backward settles it. Without model= and inputs= (or intervention=tl.options.InterventionOptions(engine="set_only") to arm in place), do refuses replay_gradient_edit_unsupported.

Hook callables receive one positional tensor and a required keyword-only ctx context. For an activation site the positional tensor is the activation; for a gradient site it is the gradient. The positional parameter is matched by position, so def hook(activation, *, ctx): ... and lambda g, *, ctx: ... are both valid. ctx.state is the per-(site, hook) scratch dict, ctx.run_state the run-scoped one, and ctx.signal is "forward" or "backward". Hooks must return the replacement tensor.

Attribution Patching¶

Attribution-patching examples in this guide use the (corrupt - clean) * patched sign convention so positive scores mean the patched activation moves the run toward the corrupt direction under that local linearization. The more common TransformerLens presentation is grad * (clean - corrupt). Both are valid as long as the sign convention is named and used consistently.

Raw PyTorch Forward Hooks¶

The TorchLens intervention API is the preferred path when you want selectors, fire records, replay, rerun, or saved intervention specs. Raw nn.Module.register_forward_hook also works during tl.trace() when you need module-local PyTorch behavior:

def halve_relu(module, args, output):
    return output * 0.5

handle = model.relu.register_forward_hook(halve_relu)
try:
    log = tl.trace(model, x)
finally:
    handle.remove()

If a raw hook returns a new tensor, TorchLens instruments that replacement, so downstream op-level graph structure remains available: torch ops the hook runs are traced as ordinary ops, and a tensor with no traceable provenance is recorded as an interventionreplacement_* op marked is_edited=True.

A capture-time module edit (intervene=tl.when(tl.module("block"), ...)) is always recorded as its own interventionreplacement_* op, with is_edited=True and its fire records on .interventions, placed between the module's last op and that op's consumers. The module's last op keeps the value it computed, before the edit. The edited value is the replacement op's .out, and it is what tl.module("block") reads as the module's output. Every later op is renumbered (a label's last number is its position), so compare traces through Bundle site reads, which pair ops by structure, not by label.

A post-capture module edit (fork.do(tl.module("block"), edit)) is recorded in the same shape: after the push the edited value sits in a new interventionreplacement_* op spliced after the module's last op, the last op keeps the value it computed, the edited module and the calls enclosing it report the replacement as their output, and a submodule whose output is that same last op keeps reporting the pre-edit value, as a plain forward hook on the outer module would leave it. Nothing is renumbered after capture: the replacement takes the next unused position number and every other op keeps its label. A later push upstream passes the parent's new value through the replacement and fires the module's hooks there again.

Comparison reads report a module edit (either kind) at the edited module's last op: tl.debug.compare and tl.debug.first_divergence, Bundle.compare, Bundle.rank_changes and the bundle diff read the replacement's value at that op's row, and the replacement op has no row of its own, so the unedited trace (which has no replacement op) is compared against the edited value at the site the edit named.

Capture-Time Recording¶

tl.record(..., save=...) is the sparse capture sibling of tl.trace(..., save=...). It returns a Recording, not a Trace; call Recording.to_trace() when you need the full postprocessed graph structure later.

recording = tl.record(model, x, save=tl.func("relu"))
trace = recording.to_trace()
streamed = tl.trace(model, x, save=tl.in_module("encoder"), storage=tl.to_disk("run.tlspec"))

record(save=...) is the only predicate spelling; the old keep_op= / keep_module= alias kwargs are removed and raise TypeError. Module-boundary event recording is gated by default_module=, which records ALL module enter/exit events uniformly — predicate-gated module-event selection has no public spelling (see docs/reference/deprecations.md for the honest capability statement).

Forward exceptions keep the historical behavior unless you opt in. With on_forward_error="attach_partial", TorchLens attaches exc.partial_recording and re-raises the original exception. With on_forward_error="return_partial", tl.record(...) returns a failed partial Recording; its output refuses with output_forward_failed because there is no valid model output. Failed partials carry status="partial_error", failed=True, string-only error metadata, num_committed_ops, and best-effort last_event_* fields. user-op failures exclude the failing call; TL-side capture failures may include a skipped/partial current-call event. Failed partials are not topology-complete, so Recording.to_trace() and Recording.backward() reject them.

A full tl.trace(...) whose forward raises raises tl.errors.TraceCaptureError, chained from the forward's own exception (exc.__cause__); inspect the partial capture on exc.partial_trace.

Trace Mutators¶

There are two distinct intervention paths, and they do not mix implicitly: (1) edit a SAVED value and push the effect downstream on the captured DAG (do() on a site or a resolved selection, direct writes), and (2) intervene on a FRESH execution (do(..., model=..., inputs=..., intervention=InterventionOptions(engine="rerun")) or a new capture with intervene=...). A new-input run_transaction(inputs=...) on a trace carrying path-1 value-edits is a fresh execution — the edits do NOT apply to it, and TorchLens discloses that at the run door with PendingValueEditsWarning (experimental name) rather than silently returning an un-edited verified run. The run itself proceeds: this is a disclosure, never a refusal.

log.do([(selection, edit), ...]) applies a list of selection and edit pairs as one transaction with one intervention event. Activation (ACT) pairs attach together and push once. A batch may also carry parameter selections (tl.param(...)) and dataflow edge selections (from log.graph.edges); on the replay engine the pairs then run in order exactly as one do() per pair would (each run of consecutive activation pairs attaches and pushes together; each parameter or edge pair goes through its own substitution path), and the result equals those calls made one at a time. The batch is atomic: if any pair is refused, every value, hook, audit row and history record the batch wrote is put back, and the trace keeps only the one error event for the failed call. Parameter and edge edits apply on the replay engine only, so a mixed batch on another engine refuses with selection_batch_cross_kind before anything runs; pass intervention=tl.options.InterventionOptions(engine="replay"). Leaf sites (inputs and buffers) still cannot join a batch (selection_batch_leaf_unsupported).

The two engines differ sharply in cost when you sweep many edits over one captured input. On the 144-head ablation sweep of a 12-layer, 12-head GPT-2 configuration (585 captured ops; benchmarks/head_sweep_engine_cost.py), the replay engine (engine="replay") took about 2.1 minutes for all 144 heads (129 s and 125 s in two runs; median 0.60-0.62 s per head) and the rerun engine (engine="rerun") about 34 minutes (2052 s and 2008 s; median 13.5-13.7 s per head), on CPU with four torch threads (torch 2.13) on a shared 32-core x86-64 Linux host. Both engines produce bitwise-identical ablated traces on all 144 heads, so the difference is cost only. A rerun re-executes every captured op, and from an intervention-ready source it also runs under the host RNG monitor, so its cost is the same for every site; replay cost grows with the edited site's downstream cone. Prefer replay for many edits on one captured input, and use rerun when the edit must see a fresh execution. A rerun from an intervention-ready source also snapshots the model state before its forward pass, as every intervention-ready capture does; each state entry whose bytes are unchanged since the source capture shares the source's snapshot instead of being cloned again, so on this sweep a rerun copies about 54 MB of tensors instead of about 706 MB, and an entry that changed (an in-place weight edit, say) is cloned afresh. Every edit on either engine starts with trace.fork(), whose cost is fixed by the size of the trace, not by the edit: on this 585-op trace a fork takes about 0.08-0.09 s and makes about 419,000 Python calls (about 580,000 before its label tables were copied directly), which is the floor under the fastest replay edits. A rerun that refreshes its fork in place releases the capture it discarded before it returns, so a fork made right after a rerun costs the same as any other (about 420,000 calls; about 1.25 million before, when that capture's cleanup ran inside whatever came next, typically the next fork).

Trace.draw(intervention_mode=...) visualizes the planned intervention recipe stored on an intervention-ready trace, such as sites registered with log.do(...). Each edited node gets a purple ring and a last row naming the edit (Δ zero_ablate), and the header names the edits; show_cone=True (off by default) adds a translucent purple underlay under the arrows and node outlines the edit changed downstream. Capture-time intervene=... calls record fire metadata on matched layers, but they are not replay plans and are not drawn as planned sites.

Bundle¶

Construct with tl.Bundle(...):

bundle = tl.Bundle({"clean": clean_log, "patched": patched_log}, baseline="clean")

Common operations:

Site reads (most_changed, compare) pair ops across members through the bundle's structural alignment, not by literal label, so an op renumbered by an inserted edit op still meets its counterpart. A site that only some members have (such as the replacement op a module edit inserts) is not scored.

Relationship gates are intentional, and they check two predicates separately: a model-axis floor (node needs same_param_shapes; comparison reads need shared_graph) and, for comparison reads, SITE LINE-UP. Runs on different inputs or checkpoints compare (clean vs. patched, two prompts, two weight states) whenever their graphs line up: the pair either captured the same input VALUES (derived from retained input payloads) or has equal graph-shape hashes. A pair on different inputs whose graphs differ (data-dependent control flow took another branch) refuses bundle_gate_input_values_differ; with no input payloads to prove identity and graphs that differ it refuses bundle_gate_input_identity_unproven. An identity relationship is never accepted as proof of input equality or of line-up. Inputs of different shapes (prompts of different token counts) refuse one of two ways: when their graphs still line up, at the first site whose shapes differ (metric_shape_mismatch, naming the site); when the length changed the graph (GPT-2 does), with bundle_gate_input_values_differ, whose message names both input shapes. Members joined by an ordering relation (successor_of/forked_from/escalates) are not comparison operands regardless of rank (bundle_gate_ordering_topology). Pair-scoped reads (compare(a, b), most_changed baseline pairs) gate only the operands they actually read.

Cohort Migration¶