transformer_lens.tools.analysis package¶

Submodules¶

Module contents¶

Analysis tools for TransformerLens.

This subpackage collects high-level, single-call interpretability analyses that sit on top of the hook/cache system. Model support is documented per tool; new analyses may target the TransformerBridge API exclusively.

Entry points and the classes they accept are exported here; result dataclasses stay in the module that returns them.

Tools:

attribution_patching: Attribution patching (gradient-linearized activation patching) over residual-stream nodes and edges: typed computational graph, a names-filtered manual-backward gradient cache, and signed node and edge scores (EAP), plus forward-only ablate-outside faithfulness with corrupt or held-out mean replacements and circuit-relative class diagnostics. Integrated gradients (EAP-IG) are not implemented.
direct_logit_attribution: Direct Logit Attribution (DLA) over components, layers, or attention heads.
get_act_patch_direct_path, get_act_patch_direct_path_all_sources: re-exported from transformer_lens.patching; exact direct path patching for head-to-head circuit analysis.
interchange: Interchange interventions on linear subspaces: set a base’s activation inside a subspace to its value on a source input, with difference-of-means, probe, random and neuron bases, a norm-matched placebo, batched gradient-free evaluation, distributed alignment search (DAS) to learn the subspace, and DAS’s trained controls (arbitrary labels and a randomly initialised model).
jacobian_lens: The Jacobian lens (J-lens) — per-layer causal transport to the output vocabulary basis, with loading of published lens artifacts, native fitting, readouts, interventions, J-space sparse decomposition, and anchored coordinate patching (offline and dynamic/hooked).
logit_lens: One-call logit lens over the accumulated residual stream and the logit_readout kernel that reads any d_model vector through the real final norm, unembedding and adapter output transform, chunked over the vocabulary with target ranks, top-k and vocabulary subsets.
projection_kernel: Basis-invariant subspace overlap and TransformerBridge attention-head OQ/OK/OV affinity.
sparse_probing: Leakage-safe k-sparse binary probes over supplied activation tensors, with train-only selection and raw null controls.

class transformer_lens.tools.analysis.EdgeAttributionConfig(granularity: Literal['node', 'edge'] = 'node', ig_steps: int = 1)¶

Bases: object

Configuration for an attribution-patching sweep.

The two axes are deliberately orthogonal:

granularity selects what is scored: "node" scores each residual-stream write, "edge" scores each (source, destination) write->read pair.
ig_steps selects gradient fidelity. ig_steps=1 is plain attribution patching / EAP: a single first-order Taylor gradient taken at the corrupt point. ig_steps>1 is EAP-IG: the integrated gradient averaged over that many points along the corrupt->clean path, which corrects the gradient saturation that makes plain attribution unfaithful.

There is intentionally no method field. An earlier design had both a method enum ("attribution"/"EAP"/"EAP-IG") and ig_steps, which overlap: the method is fully determined by granularity and whether ig_steps exceeds 1. Collapsing them removes the invalid states (e.g. method="attribution", ig_steps=5).

This build implements node and edge granularity with plain attribution. ig_steps>1 is accepted by the type but raises NotImplementedError at construction, so downstream code can import and reference this API now while the integrated-gradient path is not implemented. Once EAP-IG lands, the default flips to ig_steps=5 (EAP-IG is the faithful default); until then the default is the only executable value, ig_steps=1.

granularity¶

"node" or "edge". Defaults to "node".

Type:: Literal[‘node’, ‘edge’]

ig_steps¶

Integrated-gradient path steps (>=1). Defaults to 1.

Type:: int

granularity: Literal['node', 'edge'] = 'node'¶

ig_steps: int = 1¶

class transformer_lens.tools.analysis.FaithfulnessConfig(ablation: Literal['corrupt', 'mean'] = 'corrupt')¶

Bases: object

Configuration for an ablate-outside faithfulness measurement.

ablation¶

How an out-of-circuit edge’s writer is replaced. "corrupt" (default) substitutes the corrupt run’s own contribution; "mean" substitutes that writer’s mean over the held-out mean_tokens batch passed to faithfulness().

Type:: Literal[‘corrupt’, ‘mean’]

ablation: Literal['corrupt', 'mean'] = 'corrupt'¶

class transformer_lens.tools.analysis.JSpaceDecomposition(support: Tensor, coordinates: Tensor, selected_support: Tensor, reconstruction: Tensor, j_space_component: Tensor, non_j_space_component: Tensor)¶

Bases: object

Result of a sparse J-space decomposition.

support¶

Indices of the numerically active selected atoms – those whose nonnegative coordinate materially contributes (token ids when the dictionary is the vocabulary of J-lens vectors). A subset of selected_support.

Type:: torch.Tensor

coordinates¶

Nonnegative pursuit coefficients aligned with support (the “local J-space coordinates”); every entry is materially nonzero.

Type:: torch.Tensor

selected_support¶

Indices of every greedily selected atom, including any whose coordinate was driven to zero by the nonnegativity constraint. Defines the span for j_space_component. Satisfies support.numel() <= selected_support.numel() <= k.

Type:: torch.Tensor

reconstruction¶

The nonnegative combination sum(coordinates * active atoms) over support.

Type:: torch.Tensor

j_space_component¶

The orthogonal projection of the target onto the span of selected_support (the paper’s “J-space component”). For the exact NNLS re-solve it equals reconstruction unless a selected atom has a zero coordinate, in which case the projection uses a larger span. Numerical rank is measured on unit atom columns in float64 with relative tolerance sqrt(eps64), independently of ambient dimension and atom norms.

Type:: torch.Tensor

non_j_space_component¶

The residual target - j_space_component (the “non-J-space component”), orthogonal to the selected span.

Type:: torch.Tensor

coordinates: Tensor¶

j_space_component: Tensor¶

non_j_space_component: Tensor¶

reconstruction: Tensor¶

selected_support: Tensor¶

support: Tensor¶

class transformer_lens.tools.analysis.JacobianLens(jacobians: Dict[int, Float[Tensor, 'd_model d_model']], *, n_prompts: int, d_model: int, metadata: Dict[str, Any] | None = None)¶

Bases: object

A fitted Jacobian lens: one transport matrix per source layer.

Layer convention (matching the reference implementation and the published artifacts): index l refers to the output of block l at the Bridge-native hook blocks.{l}.hook_out. J[l] maps that activation to the final block’s output, pre final norm. The final layer itself is never fitted (its transport is the identity), so source_layers == [0, ..., n_layers - 2] for a full fit.

jacobians¶: {layer: [d_model, d_model]} transport matrices, fp32, CPU.

n_prompts¶: Number of prompts averaged into the fit.

d_model¶: Residual stream width the lens was fitted for.

metadata¶: Optional provenance (model name, TransformerLens version, fit hyperparameters). Preserved by save()/load(); artifacts from the reference implementation load with empty metadata.

ablation_hooks(model: Any, tokens: str | int | Sequence[str | int], layers: Sequence[int], *, positions: Sequence[int] | None = None) → List[Tuple[str, Any]]¶

Hooks that project token directions out of the residual stream.

With an orthonormal basis Q of the span of the tokens’ lens vectors: h <- h - Q Qᵀ h, so h ends with zero projection onto every requested vector. For a single token this is h - (h·v̂) v̂.

Parameters:

model – The model the hooks will run on.
tokens – Concept token(s) to suppress.
layers – Layers to intervene at.
positions – Chunk-local positions to ablate (negative indices allowed and normalized on every hook invocation). Defaults to all.

Returns:

[(hook_name, fn), ...] for model.hooks(fwd_hooks=...).

clear_device_cache() → None¶: Release cached Jacobians, dictionaries, and unembedding snapshots on devices.

coordinate_patch(model: Any, activation_or_prompt: Tensor | str, layer: int, source_token: str | int, target_token: str | int, *, position: int | None = None, decomposition: JSpaceDecomposition | None = None, k: int = 25, mode: str = 'substitute', alpha: float = 1.0, algorithm: str = 'nonnegative_orthogonal_matching_pursuit') → CoordinatePatch¶

Patch sparse J-space coordinates for an activation at layer.

The activation may be a raw [d_model] vector or a prompt paired with position, exactly as in decompose(). source_token must occur in the active sparse support. substitute replaces the target coordinate with the source coordinate and zeros the source; swap exchanges them. Other coordinates and x - reconstruction are fixed.

A supplied decomposition avoids repeating the vocabulary-scale sparse solve after its activation and dictionary compatibility have been validated.

Parameters:

model – A raw TransformerBridge.
activation_or_prompt – An activation vector, or a prompt (string / token tensor).
layer – Source layer (must be a fitted source layer).
source_token – Active source concept, as a single-token string or token id.
target_token – Distinct target concept, as a single-token string or token id.
position – Token position when a prompt is given; None for a raw activation.
decomposition – Optional compatible decomposition of this activation and dictionary.
k – Sparse-solver upper bound when decomposition is not supplied.
mode – "substitute" or "swap".
alpha – Finite interpolation strength; zero is an exact no-op.
algorithm – Sparse coefficient-update rule when decomposition is not supplied. A supplied decomposition is accepted if it fits as either algorithm’s output.

Returns:

A CoordinatePatch containing the edited frame, diagnostics, and activation.

coordinate_patch_hooks(model: Any, source_token: str | int, target_token: str | int, layers: Sequence[int], *, positions: Sequence[int], decomposition_cache: MutableMapping[Tuple[int, int, int], JSpaceDecomposition] | None = None, k: int = 25, mode: str = 'substitute', alpha: float = 1.0, algorithm: str = 'nonnegative_orthogonal_matching_pursuit') → List[Tuple[str, Any]]¶

Hooks that anchor-patch one J-space coordinate live, per forward-pass position.

Unlike coordinate_patch(), which edits one already-captured activation offline, this installs a forward hook that solves solve_coordinate_patch() independently for every (batch_idx, position) pair at each requested layer – a vocabulary-scale sparse decomposition per pair, per hook firing, unless decomposition_cache supplies one already validated for that (layer, batch_idx, position) key.

Parameters:

model – The model the hooks will run on.
source_token – Active source concept, as a single-token string or token id.
target_token – Distinct target concept, as a single-token string or token id.
layers – Layers to intervene at.
positions – Chunk-local positions to patch (negative indices allowed and normalized on every hook invocation). Required – there is no full-sequence default, because a silent default would trigger a vocabulary-scale solve at every position.
decomposition_cache – Optional caller-owned mapping from (layer, batch_idx, position) to a previously validated JSpaceDecomposition. A hit skips the vocabulary-scale scan; a miss solves and populates the cache. Purely a performance path – correctness does not depend on it.
k – Sparse-solver upper bound on a cache miss.
mode – "substitute" or "swap".
alpha – Finite interpolation strength; zero is an exact no-op.
algorithm – Sparse coefficient-update rule on a cache miss. A cache hit is accepted if it fits as either algorithm’s output.

Returns:

[(hook_name, fn), ...] for model.hooks(fwd_hooks=...).

Raises:

ValueError – If positions is empty, if source_token and target_token resolve to the same id, or if source_token is not in the top-k active support of every patched (batch_idx, position) pair at the moment its hook fires – the whole forward pass fails rather than silently patching a subset. This precondition is stronger and more order-dependent than “active on a clean forward pass”: in a band of layers an earlier hook’s patch edits the residual that a later layer re-decomposes, and substitute/swap zero or move the source coordinate, so the source can be removed from a later layer’s active support even though it was active on an unhooked pass. Stacking layers or positions therefore makes this progressively harder to satisfy.

Warns:

UserWarning – Once per call, naming the number of layers and positions that will perform a live vocabulary-scale solve on every cache miss.

decompose(model: Any, activation_or_prompt: Tensor | str, layer: int, *, position: int | None = None, k: int = 25, algorithm: str = 'nonnegative_orthogonal_matching_pursuit') → JSpaceDecomposition¶

Decompose an activation into its J-space content at layer.

activation_or_prompt is either:

a raw activation vector of shape [d_model] (leave position as None), or
a prompt – a string or a [1, seq] token tensor – in which case position selects the token whose blocks.{layer}.hook_out activation is decomposed.

The full-vocabulary dictionary at layer is built (and cached) via lens_vector_dictionary(), then get_sparse_decomposition() solves for a k-sparse nonnegative combination of J-lens vectors.

Parameters:

model – A raw TransformerBridge.
activation_or_prompt – An activation vector, or a prompt (string / token tensor).
layer – Source layer (must be a fitted source layer).
position – Token position when a prompt is given; must be None for a raw activation vector.
k – Upper bound on the number of J-lens vectors to select; selection stops early once no unselected vector is materially positively correlated, so fewer may be returned.
algorithm – Coefficient-update rule; see get_sparse_decomposition().

Returns:

A JSpaceDecomposition. Its support (token ids here) holds only the numerically active J-lens vectors and selected_support every selected vector, with support.numel() <= selected_support.numel() <= k. support and its token-decoding tensors are on CPU; the vector-valued outputs stay on the model’s device.

Raises:

ValueError – On an invalid model, a mismatched activation shape, a batched prompt, a prompt while the model is in training mode, an unfitted layer, or an invalid k / algorithm.
RuntimeError – If the default nonnegative least-squares solver cannot validate its result against the KKT conditions.

classmethod fit(model: Any, prompts: Sequence[str], *, corpus: str, source_layers: Sequence[int] | None = None, dim_batch: int = 8, max_seq_len: int = 128, skip_first_positions: int = 16, show_progress: bool = True, metadata: Dict[str, Any] | None = None) → JacobianLens¶

Fit a Jacobian lens on a hooked model.

Implements the reference estimator exactly. For each prompt: one forward pass (the prompt replicated dim_batch times along the batch axis), then ceil(d_model / dim_batch) backward passes. Each backward plants a one-hot cotangent for one output dimension at every valid target position simultaneously — causal attention guarantees the gradient at source position t is then the sum over target positions t' >= t with no explicit masking. Rows are averaged over valid source positions (the first skip_first_positions and the final position are excluded), and prompts contribute equally to the final mean. There is no randomness: the computation is deterministic given the prompts.

The reference implementation reports that fit quality saturates quickly — on the order of 100 prompts of 128 tokens is usable; the published lenses use up to 1000. Use merge() to parallelize across prompt slices.

Parameters:

model – A raw TransformerBridge. Model parameters are temporarily frozen (requires_grad=False) during fitting and restored after. Cotangents and activation gradients use the model dtype; fit with a float32 model for the highest-fidelity estimator. The model and all of its submodules must be in evaluation mode.
prompts – Prompt strings. Prompts too short to contain a valid position (seq_len <= skip_first_positions + 1) are skipped with a warning and do not count toward n_prompts.
corpus – Stable identifier for the prompt corpus or slice, recorded in artifact provenance.
source_layers – Layers to fit. Defaults to every layer below the final layer. Negative indices count from n_layers.
dim_batch – Output dimensions per backward pass. Higher is faster but replicates the prompt dim_batch times in memory; total backward FLOPs are unchanged.
max_seq_len – Prompts are truncated to this many tokens.
skip_first_positions – Leading positions excluded from the source average.
show_progress – Show a tqdm progress bar over prompts.
metadata – Extra provenance merged into metadata.

Returns:

The fitted JacobianLens.

Raises:

TypeError – If model is not a TransformerBridge.
ValueError – On compatibility mode, training mode, invalid provenance or layer indices, or if no prompt was long enough to fit on.

fraction_of_variance(model: Any, prompts: str | Tensor | Sequence[str | Tensor], layers: Sequence[int] | None = None, *, k: int = 25, skip_first: int = 16, positions: Sequence[int] | None = None, show_progress: bool = False) → JSpaceVarianceProfile¶

Profile the J-space share of activation variance over a prompt corpus.

Each prompt is run once (caching blocks.{layer}.hook_out for every requested layer). At each sampled position the activation is decomposed and its J-space variance fraction ||j_space_component||^2 / ||activation||^2 is recorded. The numerator is the j_space_component – the orthogonal projection of the activation onto the span of the selected support (the paper’s appendix “J-space component”), not the nonnegative reconstruction; the two coincide only when every selected atom stays active. Per layer the profile reports the median of those fractions and the pooled ratio sum(||j_space_component||^2) / sum(||activation||^2) (the paper’s “fraction of total variance”).

A layer that samples no positions – every prompt shorter than skip_first, or only zero-norm activations – contributes no fractions: its median and pooled are float("nan") and its per_position tensor is empty.

Parameters:

model – A raw TransformerBridge.
prompts – A prompt, or a sequence of prompts. Each token tensor must represent exactly one prompt and have shape [1, seq].
layers – Source layers to profile; defaults to all fitted source_layers.
k – Number of J-lens vectors per decomposition.
skip_first – Non-negative index before which positions are skipped (mirrors the fit’s early-position skip); not used for sampling when positions is given.
positions – Explicit positions to sample instead of skip_first onward.
show_progress – Show a tqdm progress bar over prompts.

Returns:

A JSpaceVarianceProfile.

Raises:

ValueError – On an invalid model, an unfitted layer, an empty corpus, a negative skip_first, a token tensor that does not have shape [1, seq], or a model in training mode.

classmethod from_pretrained(name_or_path: str, *, filename: str = 'lens.pt', revision: str | None = None, model: Any = None) → JacobianLens¶

Load a lens from a local path, a short model name, or a Hub repo.

Resolution order¶

Local file — if name_or_path is an existing .pt file, load it directly.
Local directory — if name_or_path is a directory, load <name_or_path>/<filename>.
Registry short name or HF model ID — if name_or_path matches a key or alias in the bundled jacobian_lens_registry.json (e.g. "gemma-2-2b" or "google/gemma-2-2b"), the corresponding artifact in neuronpedia/jacobian-lens is fetched automatically, at the commit the registry pins unless revision is given. The filename argument is ignored in this case because the registry already encodes the correct subpath.
Explicit Hub repo — otherwise name_or_path is treated as a Hub repo id and filename is used as-is, preserving full backward compatibility (e.g. from_pretrained("neuronpedia/jacobian-lens", filename="gpt2-small/jlens/...")).

param name_or_path:: A local .pt file, a local directory, a short model name such as "gemma-2-2b" or "llama3.1-8b", a Hugging Face model ID such as "google/gemma-2-2b", or an explicit Hub repo id paired with filename.
param filename:: File (or subpath) inside a local directory or an explicit Hub repo. Ignored when name_or_path resolves via the registry.
param revision:: Optional Hub revision (branch, tag, or commit). When omitted, registry names resolve at the registry’s pinned commit and an explicit Hub repo follows its mutable default branch; pin a commit hash for reproducible analyses of the latter.
param model:: If given, validate_model() is called so dimension or weight-processing mismatches fail here rather than at first use.
returns:: The loaded (and, if model was given, validated) lens.

Examples:

# Short model name — no need to remember the HF subpath
lens = JacobianLens.from_pretrained("gemma-2-2b", model=model)

# HF model ID also works
lens = JacobianLens.from_pretrained("google/gemma-2-2b", model=model)

# Explicit Hub repo + subpath (backward-compatible)
lens = JacobianLens.from_pretrained(
    "neuronpedia/jacobian-lens",
    filename="gpt2-small/jlens/Salesforce-wikitext/gpt2_jacobian_lens.pt",
    model=model,
)

lens_vector_dictionary(model: Any, layer: int) → Float[Tensor, 'd_vocab d_model']¶

Full-vocabulary J-lens dictionary at layer: [d_vocab, d_model].

Row t is the J-lens vector v_t = J[layer]^T W_U[:, t] – this is lens_vectors() over the entire vocabulary. The result is cached while the model’s unembedding is unchanged so a sparse decomposition can reuse it; clear_device_cache() releases it.

The dictionary is vocabulary-sized and cached on the model’s device (d_vocab * d_model fp32 values, on the order of gigabytes for a large vocabulary), one entry per requested layer. One detached copy of W_U is retained per device to detect changes without transferring weights to the host.

Parameters:

model – The model supplying W_U.
layer – Source layer for the dictionary (must be a fitted source layer).

Returns:

The dictionary, fp32, on the model’s device.

lens_vectors(model: Any, tokens: str | int | Sequence[str | int], layer: int) → Float[Tensor, 'n d_model']¶

Residual-stream directions for vocabulary tokens at a layer.

The J-lens vector for token t is row t of W_U J[layer] expressed in layer-layer residual coordinates: v_t = J[layer]^T W_U[:, t].

Parameters:

model – The model supplying W_U.
tokens – A token string / id, or a sequence of them. Strings must encode to a single token.
layer – Source layer for the vectors.

Returns:

One vector per token, fp32, on the model’s device.

classmethod load(path: str) → _LensT¶

Load a lens artifact or fit checkpoint saved in a supported schema.

Two file schemas are accepted:

Artifact (the reference format, written by save() or the Anthropic reference package): must contain a J key mapping layer indices to transport matrices, plus n_prompts, d_model, and an optional metadata dict.

Fit checkpoint (running-sum format, written by the reference implementation’s write_checkpoint() during fitting): must contain a jacobian_sum key mapping layer indices to running-sum matrices (i.e. the sum over prompts, not yet divided by the prompt count), plus n_done. d_model is inferred from the first matrix’s shape; no explicit d_model key is required or expected. The per-layer means are reconstructed on load. A converted_from: "jacobian_lens_checkpoint" key is added to metadata so merge() refuses to silently combine checkpoints with natively TL-fitted lenses. Fit-reserved provenance keys (transformer_lens_fit, etc.) are stripped; scalar fields that can be serialised under weights_only=True are preserved. Tensor-valued metadata fields that would fail _validate_metadata() are recorded by name in a dropped_fields list.

Fit checkpoint schema (reference `write_checkpoint()` format)¶

The reference implementation writes exactly six top-level keys; all other keys in the payload are ignored:

{
    "jacobian_sum":   {<layer int>: <float32 tensor [d, d]>, ...},
    "n_done":         <int>,        # prompts accumulated into jacobian_sum
    "next_idx":       <int>,        # next prompt index (informational)
    "source_layers":  [<int>, ...], # documented layer indices (informational)
    "target_layer":   <int>,        # target layer — harvested into metadata
    "skip_first":     <int>,        # leading positions skipped (informational)
    # optional flat provenance accepted from alternative checkpoint writers:
    "model_name":     <str>,
    "model_revision": <str>,
    "corpus":         <str>,
    # optional nested provenance accepted from alternative writers:
    "metadata":       {<str>: <scalar/list/dict>, ...},
}

param path:: Path to the .pt file.
raises ValueError:: If the file lacks both a J key (artifact) and a jacobian_sum key (checkpoint), or if a checkpoint records a non-positive n_prompts.

classmethod merge(lenses: Sequence[JacobianLens]) → JacobianLens¶

Combine lenses fitted on disjoint prompt slices.

The per-layer matrices are averaged weighted by each lens’s n_prompts, matching the reference implementation, so fitting can be parallelized across processes or machines and merged afterwards. Provenance must match across shards (apart from n_prompts), so a merge cannot silently relabel matrices fitted with different models, corpora, dtypes, or recorded estimator settings. The merged count replaces the per-shard count.

Parameters:: lenses – Lenses that agree exactly on source_layers and d_model.
Raises:: ValueError – On an empty sequence, mismatched lenses, or shards from different estimators.

occupancy(model: Any, activation_or_prompt: Tensor | str, layer: int, *, position: int | None = None, max_atoms: int = 25, num_control_dictionaries: int = 32, seed: int = 0) → JSpaceOccupancy¶

Estimate how many J-lens vectors are meaningfully active in an activation at layer.

Resolves activation_or_prompt (a raw [d_model] vector, or a prompt plus position) exactly as decompose(), builds the cached full-vocabulary dictionary via lens_vector_dictionary(), and calls estimate_occupancy().

Control dictionaries are [d_vocab, d_model] CPU draws (for cross-device reproducibility) that dominate run time on an accelerator; lower num_control_dictionaries for speed.

Parameters:

model – A raw TransformerBridge.
activation_or_prompt – An activation vector, or a prompt (string / token tensor).
layer – Source layer (must be a fitted source layer).
position – Token position when a prompt is given; None for a raw activation.
max_atoms – Maximum number of J-lens vectors to consider.
num_control_dictionaries – Number of random control dictionaries to average over.
seed – Seed for the random control dictionaries (reproducibility).

Returns:

A JSpaceOccupancy.

readout(model: Any, input: str | Int[Tensor, 'batch seq'], *, layers: Sequence[int] | None = None, positions: Sequence[int] | None = None, use_jacobian: bool = True, top_k: int = 10, return_full_logits: bool = False) → JacobianLensReadout¶

Read per-layer vocabulary logits for a prompt.

Runs the model once with caching, transports the residual stream at each requested layer through J[layer] (or the identity when use_jacobian=False — the logit lens), and applies the model’s own final norm, unembedding, architecture logit scaling, and logit soft cap.

Parameters:

model – A raw TransformerBridge.
input – A prompt string, or a [1, seq] token tensor.
layers – Layers to read. Defaults to every fitted layer plus the final layer. The final layer (n_layers - 1) is always read with the identity transport — by construction its lens equals the model’s own output distribution.
positions – Token positions to read (negative indices allowed). Defaults to all positions.
use_jacobian – Apply the Jacobian transport. False gives the logit-lens baseline through the identical code path.
top_k – Number of values and vocabulary ids retained per layer and position. Defaults to 10.
return_full_logits – Also retain full vocabulary tensors on CPU. This is opt-in because a 64-token Gemma readout across all layers is roughly 1.7 GB.

Returns:

A JacobianLensReadout.

Raises:

ValueError – If the model fails validate_model(), input is batched, top_k is invalid, a requested layer has no transport matrix, or the model is in training mode.

save(path: str, *, dtype: dtype = torch.float16) → None¶

Save the lens in the reference implementation’s artifact format.

The four official keys (J, n_prompts, source_layers, d_model) are written unchanged so the file stays loadable by the reference package; TransformerLens provenance is stored under an additive metadata key.

Parameters:

path – Destination .pt path.
dtype – Storage dtype. Defaults to fp16 like the reference implementation — Jacobian entries are order-one, so the smaller dtype costs little precision and halves the artifact on disk.

property source_layers: List[int]¶: Sorted list of layers this lens has transport matrices for.

steering_hooks(model: Any, token: str | int, layers: Sequence[int], *, alpha: float = 4.0, positions: Sequence[int] | None = None) → List[Tuple[str, Any]]¶

Hooks that steer the residual stream along a token’s J-lens vector.

At each layer the unit-normalized lens vector is added, scaled by alpha times the activation’s median per-position residual norm: h <- h + alpha * median||h|| * v̂. This norm-matched parameterization follows the steering description in the reference implementation’s experiment protocols; the paper’s minimal form is the unscaled h <- h + alpha * v_t, recoverable by passing the raw lens_vectors() output to your own hook. The median (not mean) is used so attention-sink positions — whose residual norms run orders of magnitude above typical positions — do not inflate the scale.

Parameters:

model – The model the hooks will run on.
token – The concept token to steer toward.
layers – Layers to intervene at.
alpha – Steering strength scalar; 0 disables. Because of the norm-matched scale, values of order 1 already perturb the stream by roughly its own magnitude.
positions – Chunk-local positions to steer (negative indices allowed and normalized on every hook invocation). Defaults to all.

Returns:

[(hook_name, fn), ...] for model.hooks(fwd_hooks=...) or model.run_with_hooks(fwd_hooks=...).

swap_clamp_hooks(model: Any, source_token: str | int, target_token: str | int, layers: Sequence[int], clean_cache: ActivationCache | Mapping[str, Tensor], *, positions: Sequence[int] | None = None) → List[Tuple[str, Any]]¶: Deprecated: use swap_hooks() with clamp_to=clean_cache.

swap_hooks(model: Any, source_token: str | int, target_token: str | int, layers: Sequence[int], *, alpha: float = 1.0, positions: Sequence[int] | None = None, clamp_to: ActivationCache | Mapping[str, Tensor] | None = None) → List[Tuple[str, Any]]¶

Hooks that exchange two concepts’ coordinates in lens space.

The paper’s patching-in-lens-coordinates intervention: with V = [v_s, v_t] and lens coordinates c = V⁺ h (pseudoinverse), the update is h <- h + alpha * V (c_target - c). The component of h orthogonal to span{v_s, v_t} is untouched. alpha=2 is the paper’s “double-strength” swap.

Without clamp_to, c_target = sigma(c) exchanges the two live coordinates, re-read from the activation seen by every hook. It is therefore an involution when applied repeatedly in a subspace whose coordinates are preserved between layers: a second application can undo the first.

With clamp_to (the paper’s multi-layer clamp protocol), c_target is the exchange of each layer’s clean-run coordinates, fixed when the hooks are built, so at alpha=1 the update is idempotent at each layer: h <- h + V (c_target - V⁺h).

Parameters:

model – The model the hooks will run on.
source_token – The concept to remove (e.g. " France").
target_token – The concept to install (e.g. " China").
layers – Layers to intervene at.
alpha – Swap strength; with clamp_to, the fraction of the way to the clamped coordinates.
positions – Chunk-local positions to swap (negative indices allowed and normalized on every hook invocation, and against the clean activations when clamping). Defaults to all.
clamp_to – Activations from an unmodified run_with_cache at each requested layer’s blocks.{layer}.hook_out name. Either the ActivationCache it returns by default or the plain dict from return_cache_object=False is accepted.

Returns:

[(hook_name, fn), ...] for model.hooks(fwd_hooks=...).

transport(residual: Float[Tensor, '... d_model'], layer: int) → Float[Tensor, '... d_model']¶

Map layer-layer activations into the final block’s output basis.

Computes J[layer] @ h per activation vector, in fp32.

Parameters:

residual – Activations from the output of block layer.
layer – Source layer index.

validate_model(model: Any) → JacobianLens¶

Check that model matches this lens; raise loudly if not.

Requires a raw causal TransformerBridge with the standard direct final-norm/unembed path, verifies recorded model provenance, residual width and layer range, and enforces the published final-block target convention.

Parameters:

model – A raw TransformerBridge.

Returns:

self, for chaining.

Raises:

TypeError – If model is not a TransformerBridge.
ValueError – On model provenance or d_model mismatch, out-of-range source layers, compatibility mode, unsupported attention/output paths, or a non-final target convention.

class transformer_lens.tools.analysis.LearnedSubspace(d_model: int, rank: int)¶

Bases: Module

A trainable rank-rank subspace of d_model whose basis stays orthonormal.

Parameters:: rank – Number of basis directions, in [1, d_model].

property basis: Float[Tensor, 'd_model rank']¶

class transformer_lens.tools.analysis.Node(kind: Literal['embed', 'attn_head_out', 'mlp_out', 'q_input', 'k_input', 'v_input', 'mlp_in', 'logits'], position: int, layer: int | None = None, head: int | None = None)¶

Bases: object

A node in the residual-stream computational graph.

Nodes are the typed, hashable keys the attribution sweep scores. Each node is identified by (kind, layer, position, head); kind selects the node family and constrains which of layer/head apply. Three kinds are writers – they contribute a value into the residual stream:

"embed": the token embedding write. layer and head are None.
"attn_head_out": one attention head’s output. layer and head set.
"mlp_out": one layer’s MLP output. layer set, head is None.

Five kinds are readers – they consume the residual stream as an edge’s destination (see enumerate_edges()):

"q_input" / "k_input" / "v_input": one attention head’s split Q/K/V input. layer and head set.
"mlp_in": one layer’s MLP entry. layer set, head is None.
"logits": the terminal readout of the final residual, read at blocks.{n_layers-1}.hook_resid_post. layer is that final layer and head is None. Every writer feeds this reader, so a writer’s aggregate over its outgoing edges equals its direct node score.

position is the sequence index the node is read at. The invariants above are enforced in __post_init__ so a malformed key raises rather than silently producing a wrong graph.

head: int | None = None¶

property hook_name: str¶

The cache hook point this node reads from.

Uses the standard TransformerBridge alias names (hook_embed, blocks.{l}.attn.hook_z, blocks.{l}.hook_mlp_out, blocks.{l}.attn.hook_q_input/hook_k_input/hook_v_input, blocks.{l}.hook_mlp_in, blocks.{l}.hook_resid_post); the per-head nodes (attn_head_out, q_input, k_input, v_input) slice head self.head out of the shared per-head tensor.

kind: Literal['embed', 'attn_head_out', 'mlp_out', 'q_input', 'k_input', 'v_input', 'mlp_in', 'logits']¶

layer: int | None = None¶

position: int¶

class transformer_lens.tools.analysis.NormMatchedPlacebo(reference: SubspaceInterchange, basis: Tensor)¶

Bases: object

Move each base exactly as far as reference does, but inside a random subspace.

An effect that the placebo also produces comes from displacing the activation by that much, not from the aligned subspace.

reference¶

The intervention whose displacement is matched.

Type:: transformer_lens.tools.analysis.interchange.SubspaceInterchange

basis¶

Orthonormal basis of the placebo subspace, usually from random_basis() with the reference’s rank, or from random_neuron_basis() for a reference built by neuron_basis().

Type:: torch.Tensor

basis: Tensor¶

reference: SubspaceInterchange¶

exception transformer_lens.tools.analysis.SparseProbeConvergenceError¶

Bases: RuntimeError

A probe fit failed to converge or produced non-finite output.

Distinct from other ``RuntimeError``s so the sweep’s control-rejection path swallows only convergence failures, never unrelated runtime faults.

class transformer_lens.tools.analysis.SubspaceBasis(basis: Float[Tensor, 'ambient rank'], singular_values: Float[Tensor, 'spectrum'], rank: int, measured_rank: int, rtol: float, threshold: float, input_shape: Tuple[int, int])¶

Bases: object

An explicitly ranked orthonormal basis extracted from a matrix.

basis¶

Orthonormal column-space basis, [ambient_dim, rank].

Type:: jaxtyping.Float[Tensor, ‘ambient rank’]

singular_values¶

All reduced-SVD singular values, in descending order.

Type:: jaxtyping.Float[Tensor, ‘spectrum’]

rank¶

Number of retained basis directions.

Type:: int

measured_rank¶

Numerical rank before optional caller truncation.

Type:: int

rtol¶

Effective relative rank tolerance.

Type:: float

threshold¶

Absolute singular-value threshold used for rank measurement.

Type:: float

input_shape¶

Shape of the matrix from which the basis was extracted.

Type:: Tuple[int, int]

property ambient_dim: int¶: Dimension of the space containing the subspace.

basis: Float[Tensor, 'ambient rank']¶

input_shape: Tuple[int, int]¶

measured_rank: int¶

rank: int¶

rtol: float¶

singular_values: Float[Tensor, 'spectrum']¶

threshold: float¶

class transformer_lens.tools.analysis.SubspaceInterchange(basis: Tensor | LearnedSubspace, strength: float = 1.0)¶

Bases: object

Set the base’s coordinates inside basis to the source’s.

basis¶

Orthonormal basis, [d_model, rank], or a LearnedSubspace.

Type:: torch.Tensor | transformer_lens.tools.analysis.interchange.LearnedSubspace

strength¶

1 interchanges the coordinates; other values scale the move, 0 leaves the base.

Type:: float

basis: Tensor | LearnedSubspace¶

strength: float = 1.0¶

transformer_lens.tools.analysis.arbitrary_targets(keys: Sequence[Hashable], n_classes: int, seed: int) → Int[Tensor, 'batch']¶

A fixed pseudo-random class per key: targets for the arbitrary-label control.

Train fit_das() on these instead of the high-level model’s targets, with the same pairs, rank and budget. The interchange accuracy it reaches shows how much accuracy the method’s expressivity alone produces; a learned subspace should clearly exceed it.

Parameters:: keys – One key per pair identifying its source, e.g. the source text.
Returns:: A class in [0, n_classes) per key; equal keys always get equal classes.

transformer_lens.tools.analysis.attention_head_subspace_affinity(model: Any, *, source_role: str = 'O', target_role: str, layer_order: str = 'forward', rank: int | None = None, rtol: float | None = None) → HeadAffinityResult¶

Compute OQ, OK, or OV Projection Kernel affinities for a TransformerBridge.

K/V axes preserve native key-value heads on grouped-query attention models; they are never expanded to query-head count. Hybrid models include only attention blocks and report their original block indices.

Parameters:

model – A TransformerBridge exposing readable per-head attention weights.
source_role – Source role; v1 supports only "O".
target_role – One of "Q", "K", or "V".
layer_order – "forward" keeps strict earlier-to-later pairs; "all" keeps every pair.
rank – Optional common truncation rank. By default every head must be full column rank.
rtol – Optional relative numerical-rank tolerance.

Returns:

Affinity tensors, validity mask, original layer indices, and rank metadata.

transformer_lens.tools.analysis.attribution_patch(model: Any, clean: Tensor, corrupt: Tensor, metric_fn: Callable[[Tensor], Tensor], config: EdgeAttributionConfig = EdgeAttributionConfig(granularity='node', ig_steps=1)) → AttributionResult¶

Estimate every component’s causal effect on metric_fn in two forwards + one backward.

For each clean/corrupt pair this runs a clean forward (for a_clean) and a corrupt forward whose backward hooks capture g = d(metric)/d(a) (for a_corrupt and its gradient). At node granularity (config.granularity == "node") each node is scored with the first-order Taylor estimate effect(node) = (a_clean - a_corrupt) . g. At edge granularity (config.granularity == "edge") each writer -> reader edge is scored with effect(edge) = (a_clean[writer] - a_corrupt[writer]) . d(metric)/d(input of reader), and node_scores holds each writer’s aggregate effect (the sum of its outgoing edge scores).

Sign/direction convention (denoising form): gradients are taken on the corrupt run and the estimate points toward the clean activation, so a positive score means patching that node from corrupt toward clean moves the metric in the positive direction. An oracle-parity test maps this convention onto a pinned reference rather than assuming the two agree.

Dataset averaging: clean/corrupt may hold a batch of prompt pairs. Each pair is scored independently (per-example forward/backward, so its own reconstruction identity holds) and per-node (or per-edge) scores are averaged across the batch before ranking.

The model and every submodule must be in evaluation mode. Separate clean and corrupt forwards cannot produce meaningful activation differences if stochastic training layers such as dropout remain active.

Parameters:

model – A TransformerBridge (or compatible) exposing cfg.n_layers, hook_dict, and hooks().
clean – Clean token ids, shape [batch, seq].
corrupt – Corrupt token ids, shape [batch, seq], paired row-by-row with clean.
metric_fn – Maps single-example logits to a scalar to differentiate.
config – Sweep configuration. Node and edge granularity are both supported with plain attribution (ig_steps=1); ig_steps>1 raises at construction.

Returns:

An AttributionResult whose node_scores are averaged over the batch. For an edge-granularity sweep, edge_scores is populated too and node_scores is the per-writer aggregate of those edge scores. The graph’s terminal logits reader gives every writer an edge carrying its direct skip-connection contribution to the metric, so this aggregate equals the quantity a node-granularity sweep on the same model returns (including for the final layer’s MLP output, whose only outgoing edge is the one to the logits reader).

Raises:

ValueError – if clean/corrupt are not 2D, hold a different number of pairs, a pair tokenizes to different lengths (activations must align position-by-position), or the model or one of its submodules is in training mode. At edge granularity, also if per-head attention outputs are normalized or scaled before the residual add (post-norm, scaled-residual or parallel-mixer blocks), where attn.hook_result is not a head’s residual contribution, or if an attention layer’s K/V head count differs from n_heads (GQA/MQA).

transformer_lens.tools.analysis.collect_activations(model: Any, hook_name: str, inputs: Tensor, task_logits: Callable[[Tensor], Float[Tensor, 'batch n_classes']], *, positions: Int[Tensor, 'batch'] | None = None, batch_size: int = 256) → Float[Tensor, 'batch d_model']¶

Activations at hook_name, at one position per example if given, returned on the CPU.

Parameters:

model – A model exposing the hook API.
inputs – Model inputs, one row per example.
task_logits – Runs the model on a batch of inputs.
positions – For [batch, pos, d_model] activations, the position to read per example.
batch_size – Rows per forward pass.

Returns:

Activations, [batch, d_model].

Raises:

ValueError – If any part of the model is in training mode, or the activation does not match the positions given (or not given).
RuntimeError – If the hook does not fire exactly once per forward pass.

Compute Direct Logit Attribution for a prompt.

Decomposes the contribution of model components to the logit of answer_tokens (or, if incorrect_tokens is given, to the logit difference answer - incorrect along the W_U direction, which is usually what you want for circuit analysis).

The model is run once with caching unless a precomputed cache is passed.

Note that DLA attributes only the part of a logit that comes from the residual stream through the unembedding direction; the unembedding bias b_U is a per-token constant that no component produces. So a complete decomposition reconstructs logit[token] - b_U[token] rather than the raw logit.

On a TransformerBridge, compatibility mode must be enabled and the final normalization’s affine parameters must be folded into W_U and b_U (or already be identity); otherwise the projection direction is wrong and DLA would return silently incorrect numbers. Hybrid architectures (Mamba/SSM/Mixer/LinearAttention) are not yet supported because decompose_resid only understands the attn_out + mlp_out block layout; both conditions raise an explicit error at call time, as does a Bridge whose readout is not ln_final -> unembed on the last block’s output; a prompt run here is also checked against the model’s own output, a precomputed cache only structurally.

Parameters:

model – A TransformerBridge with enable_compatibility_mode() already called and an identity final normalization.
input – Prompt to run — a string, list of strings, or token tensor. Optional only when a precomputed cache is supplied.
answer_tokens – The correct token(s) to attribute, as a string, id, or tensor. A string is converted with model.to_single_token.
incorrect_tokens – Optional baseline token(s). When given, attribution is computed for the answer - incorrect residual direction. Must broadcast to the same shape as answer_tokens.
unit –
Decomposition granularity:
- "component" (default): embedding + each layer’s attention and MLP output (via decompose_resid).
- "layer": cumulative residual stream after each sublayer, i.e. logit-lens trajectory (via accumulated_resid).
- "head": each attention head individually, plus a remainder term for everything else (via stack_head_results).
pos – Sequence position(s) to attribute. Defaults to -1 (the final token, the usual choice for next-token DLA). Pass None to keep every position (the result then has a trailing position axis).
cache – Optional precomputed ActivationCache to reuse instead of running the model again.

Returns:

A DirectLogitAttribution with attribution (shape [component, *batch_and_pos]) and aligned labels.

Raises:

ValueError – If unit is invalid, answer_tokens is None, neither input nor cache is provided, or a TransformerBridge is passed without compatibility mode enabled, with an active learned final-norm weight or bias, or with a readout that is not ln_final -> unembed; also if the model must run while any module is in training mode (a supplied cache is not checked).
NotImplementedError – If a TransformerBridge reports a hybrid block layout (Mamba/SSM/Mixer/LinearAttention).

transformer_lens.tools.analysis.faithfulness(model: Any, clean: Tensor, corrupt: Tensor, metric_fn: Callable[[Tensor], Tensor], circuit: Sequence[tuple[Node, Node]] | Sequence[tuple[Node, Node, float]], config: FaithfulnessConfig = FaithfulnessConfig(ablation='corrupt'), *, mean_tokens: Tensor | None = None) → FaithfulnessResult¶

Measure how much of the clean-to-corrupt metric gap a circuit recovers.

Ablates every edge outside circuit (forward-only, exact rather than linearized) and reports the metric the model then produces. Feed a ranked edge set straight in from attribution_patch(..., config=EdgeAttributionConfig(granularity="edge")).top_edges(k=...).

Under mean ablation recovery is still normalized to the clean/corrupt gap, so an empty circuit need not have zero recovery.

Parameters:

model – A TransformerBridge (or compatible) exposing cfg.n_layers, hook_dict, and hooks().
clean – Clean token ids for one evaluated example, shape [1, seq].
corrupt – Paired corrupt token ids, shape [1, seq].
metric_fn – Maps single-example logits to a scalar.
circuit – The edges to keep, as (writer, reader) pairs or the (writer, reader, score) triples top_edges returns. Every other edge in the graph is ablated. Duplicate entries count once.
config – Ablation configuration. Defaults to replacing an ablated writer with the corrupt run’s own contribution.
mean_tokens – Required only for ablation="mean". Held-out int32/int64 calibration token ids, [n_calibration, seq] at the evaluated length and on the clean/corrupt device; writer activations are averaged per position and head. Rows equal to either endpoint are rejected; avoid semantic leakage too.

Returns:

A FaithfulnessResult. The per-edge-class breakout costs one extra forward per retained class.

Raises:

ValueError – if clean/corrupt are not 2D, hold a different number of pairs, hold more than one pair, a pair tokenizes to different lengths, the model or a submodule is in training mode, a circuit edge is not in the graph, the clean and corrupt metrics are equal (so the recovered fraction is undefined), mean calibration is missing, empty, misaligned, on a different input device, has an unsupported dtype, contains an evaluated endpoint, or is supplied in corrupt mode; or per-head attention outputs are normalized or scaled before the residual add (post-norm, scaled-residual or parallel-mixer blocks), or an attention layer’s K/V head count differs from n_heads (GQA/MQA).

transformer_lens.tools.analysis.fit_das(model: Any, hook_name: str, base_inputs: Tensor, source_inputs: Tensor, targets: Int[Tensor, 'batch'], task_logits: Callable[[Tensor], Float[Tensor, 'batch n_classes']], *, rank: int, base_positions: Int[Tensor, 'batch'] | None = None, source_positions: Int[Tensor, 'batch'] | None = None, epochs: int = 20, batch_size: int = 128, lr: float = 0.01, strength: float = 1.0, seed: int = 0, device: str | device | None = None, micro_batch_size: int | None = None) → DASResult¶

Learn the subspace whose interchange makes the model output targets (DAS).

targets[i] is the class the high-level causal model predicts for base i with the aligned variable set to its value in source i. The basis is trained with cross-entropy on the task logits under the interchange. Model parameters are frozen during the call and their requires_grad flags restored afterwards.

Train on pairs where several variables differ between base and source, labelled by the high-level model’s interchange, as well as on pairs that differ only in the aligned variable: otherwise the subspace can learn whatever the pairs happen to move together.

Parameters:

model – A model exposing the hook API.
hook_name – The hook point to align.
base_inputs – Base inputs, one row per pair.
source_inputs – Source inputs, one row per pair.
targets – The high-level model’s output class for each pair under the interchange.
task_logits – Runs the model on a batch of inputs and returns task logits.
rank – Dimension of the subspace.
base_positions – For [batch, pos, d_model] activations, the base position per pair.
source_positions – Likewise for the sources.
epochs – Passes over the pairs.
batch_size – Pairs per optimisation step.
lr – Adam learning rate.
strength – Interchange strength during training.
seed – Seed for the basis initialisation and the batch order.
device – Device for the subspace’s parameters; by default the device of the model’s parameters, or the CPU for a model without any.
micro_batch_size – Split each batch into parts this size and accumulate gradients: the same gradients, up to floating-point rounding, in less memory.

Returns:

The learned subspace with its training losses.

Raises:

ValueError – If any part of the model is in training mode, or the inputs and targets differ in length.

transformer_lens.tools.analysis.fit_sparse_probe(features: Float[Tensor, 'example feature'], labels: Bool[Tensor, 'example'] | Integer[Tensor, 'example'], *, k: int, test_fraction: int | float = 0.3, positive_label: int | bool = 1, groups: Integer[Tensor, 'example'] | None = None, preprocess: str = 'none', std_floor: int | float = 0.001, class_weight: str | None = 'balanced', l2_strength: int | float = 0.01, seed: int = 0, max_iter: int = 200, max_refinement_steps: int = 25, decrement_tolerance: int | float = 1e-12) → SparseProbeResult¶

Fit a train-only-selected k-sparse binary logistic probe.

The stratified split is created before feature scoring or optional standardization. Floating result tensors are detached CPU float64 tensors.

Parameters:

features – Finite float16/bfloat16/float32/float64 tensor shaped [example, feature].
labels – Boolean or integer binary labels shaped [example].
k – Number of coordinates selected by absolute train class-mean difference.
test_fraction – Requested held-out fraction within each class. With groups it targets held-out rows per class, which whole groups only approximate.
positive_label – Label defining the positive class and score sign.
groups – Optional integer group IDs; each group is held out whole, with both classes on both sides. Groups may mix classes, but each class must appear in two or more.
preprocess – "none" or train-only "standardize".
std_floor – Lower bound on a non-constant column’s "standardize" scale, so a near-constant column is not amplified to unit scale; 0 disables it.
class_weight – "balanced" or None for unweighted BCE.
l2_strength – Positive coefficient penalty in the logistic objective.
seed – Local CPU-generator seed used only for the stratified split.
max_iter – Maximum LBFGS iterations before Newton refinement. LBFGS also stops once its function evaluations reach max_iter * 5 // 4 (stop_reason="max_eval"; the line search can finish one past it), a budget equal to max_iter when max_iter <= 3.
max_refinement_steps – Maximum damped Newton steps after LBFGS; zero only checks.
decrement_tolerance – Largest accepted Newton decrement g^T H^-1 g / 2, an estimate of the objective gap to the optimum in nats; a larger gap raises.

Returns:

Selected support, fitted parameters, split/preprocessing metadata, metrics, and optimizer diagnostics.

Raises:

ValueError – If inputs, groups, or options violate the binary-probe contract.
SparseProbeConvergenceError – If the optimizer fails or misses its convergence threshold.

transformer_lens.tools.analysis.get_act_patch_direct_path(model: TransformerLensModel, corrupted_tokens: Tensor, clean_cache: ActivationCache, corrupted_cache: ActivationCache, patching_metric: Callable[[Tensor], Tensor], src_layer: int, src_head: int, component: Literal['q', 'k', 'v'] = 'q', verbose: bool = True) → Float[Tensor, 'n_layers n_heads']¶

Patch one source head’s output into each later head’s Q, K or V input alone.

Direct path patching (Wang et al. 2022) isolates the edge A -> B: A’s clean-minus-corrupted write (z @ W_O) is added at every position to B’s own pre-ln1 copy of the residual, while every other reader still sees A’s corrupted output. The result is the exact edge intervention, folded LayerNorm or not.

The sweep temporarily enables cfg.use_split_qkv_input (restoring the caller’s flags on exit), which costs about 3 * n_heads times the usual ln1 activation memory per forward.

Parameters:

model – A TransformerBridge whose attention exposes the pre-Q/K/V fork.
corrupted_tokens – Token ids for the corrupted run, [batch, pos].
clean_cache – Cache of the clean run; must hold blocks.{src_layer}.attn.hook_z.
corrupted_cache – Cache of the corrupted run; must hold the same hook_z.
patching_metric – Maps the patched logits to a scalar.
src_layer – Layer of the source head.
src_head – Index of the source head.
component – Destination input to patch: "q" (default), "k" or "v".

Returns:

results[l, h] is the metric with the path into head (l, h) patched. Rows l <= src_layer are 0.0 placeholders, not measurements.

Return type:

[n_layers, n_heads]

Raises:

ValueError – if a block from src_layer on has no attention, if the source block normalises, scales or joins its attention output before the residual add (post-norm, scaled-residual or parallel-mixer blocks), if component is "k" or "v" on a GQA/MQA model, or if any module is in training mode.
NotImplementedError – if a destination layer’s attention has no pre-Q/K/V fork.

transformer_lens.tools.analysis.get_act_patch_direct_path_all_sources(model: TransformerLensModel, corrupted_tokens: Tensor, clean_cache: ActivationCache, corrupted_cache: ActivationCache, patching_metric: Callable[[Tensor], Tensor], component: Literal['q', 'k', 'v'] = 'q', verbose: bool = True) → Float[Tensor, 'n_layers n_heads n_layers n_heads']¶

Run get_act_patch_direct_path() from every source head.

results[sl, sh, dl, dh] is the metric with the path from head (sl, sh) into head (dl, dh) patched, and 0.0 where dl <= sl. One forward per (source, later destination) pair, so it suits small models or targeted sub-sweeps.

transformer_lens.tools.analysis.get_sparse_decomposition(x: Tensor, dictionary: Tensor, k: int = 25, *, algorithm: str = 'nonnegative_orthogonal_matching_pursuit') → JSpaceDecomposition¶

Greedily decompose x into a k-sparse nonnegative combination of atoms.

Parameters:

x – Target vector, shape [d_model].
dictionary – Atom matrix, shape [num_atoms, d_model] (rows are atoms).
k – Upper bound on the number of atoms to select. Selection stops early once no unselected atom is materially positively correlated with the residual, so fewer than k atoms may be selected (and fewer still may be numerically active).
algorithm – Coefficient-update rule. "nonnegative_orthogonal_matching_pursuit" (default) re-solves the selected-set coefficients exactly as a nonnegative least-squares fit; "gradient_pursuit" takes a single projected-gradient step per atom. See the module docstring for the trade-off (they use the same selection rule, while the exact re-solve is optimal on each selected set).

Returns:

A JSpaceDecomposition. Its support holds only the numerically active atoms and selected_support every selected atom, with support.numel() <= selected_support.numel() <= k.

Raises:

ValueError – On an unknown algorithm, complex or non-finite inputs, a non-2-D dictionary, a target whose length does not match d_model, k outside [1, num_atoms], or a dictionary with non-finite or zero-norm atoms.
RuntimeError – If algorithm="nonnegative_orthogonal_matching_pursuit" and the nonnegative least-squares solve cannot be certified against its KKT conditions within its numerical tolerance.

transformer_lens.tools.analysis.intervened_logits(model: Any, hook_name: str, intervention: Callable[[Tensor, Tensor], Tensor], base_inputs: Tensor, source_activations: Float[Tensor, 'batch d_model'], task_logits: Callable[[Tensor], Float[Tensor, 'batch n_classes']], *, positions: Int[Tensor, 'batch'] | None = None, batch_size: int = 256) → Float[Tensor, 'batch n_classes']¶

Task logits under intervention, in batches without gradients, returned on the CPU.

The evaluation counterpart of run_with_intervention(), with the same arguments and checks: memory is bounded by batch_size rather than by the number of pairs.

Parameters:

model – A model exposing the hook API.
intervention – Maps base and source activations to new base activations.
source_activations – Activations from the source inputs at the same hook point, one row per base.
task_logits – Runs the model on a batch of inputs.
positions – For [batch, pos, d_model] activations, the position to intervene on per base.
batch_size – Rows per forward pass.

Returns:

Task logits under the intervention, [batch, n_classes].

Raises:

ValueError – If any part of the model is in training mode, or the inputs do not pair up.

Read the accumulated residual stream through the final norm and unembedding.

Builds the stack with accumulated_resid(), optionally applies a per-entry lens translator, applies the final norm through apply_ln_to_stack() with recomputed statistics, and hands the result to logit_readout(). The final stack entry reproduces the model’s logits (raw bridge) or log-probabilities (compatibility mode). Bridges whose readout is not ln_final -> unembed on the last block’s output are refused; a prompt run here with apply_ln=True is also checked against the model’s own output, a supplied cache only structurally.

Parameters:

model – A TransformerBridge, or an ActivationCache (then input must be omitted).
input – Prompt(s) or token tensor to run when no cache is given. Only the residual-stream hooks are cached. A batch of strings is padded by the tokenizer, so with the default positions=-1 the shorter prompts are read at a pad position; pass a token tensor (or per-prompt positions) when prompts differ in length.
cache – An existing cache to read instead of running the model.
layers – Indices into the accumulated stack (negative allowed); None keeps every entry.
incl_mid – Include resid_mid entries. Raises on hybrid stacks whose SSM blocks have no resid_mid.
positions – Position selection (Slice semantics); default last token.
batch_slice – Batch selection (Slice semantics).
lens – Optional f(entry, index) -> entry applied to each stack entry before the norm, e.g. a tuned-lens translator. Combine with enable_grad=True to fit or attribute through it.
targets – As in logit_readout(); a per-row tensor is aligned with the stack’s [batch, pos] dims after positions and batch_slice are applied ([batch, pos] next-token ids with positions=None, [batch] with the default last position).
apply_ln – As in logit_readout().
vocab – As in logit_readout().
return_type – As in logit_readout().
top_k – As in logit_readout().
chunk_size – As in logit_readout().
enable_grad – As in logit_readout().

Returns:

A LogitLensResult.

Read d_model vectors through the model’s final norm and unembedding.

Computes transform(ln_final(x) @ W_U + b_U) for every vector, where transform is the adapter’s post-unembedding transform, in chunks of chunk_size rows so at most chunk_size × d_vocab logits exist at once. Everything derived from the full vocabulary (the softmax normalizer, entropy, target_ranks) is computed per chunk and discarded; only the selected columns are kept.

Parameters:

model – A TransformerBridge in eval mode.
vectors – Residual-stream-shaped vectors, any leading dims.
apply_ln – Apply the model’s ln_final module first (recomputed statistics per vector). False reads the raw vectors and does not reproduce the model.
vocab – Token ids to keep along the trailing axis.
targets – Tokens whose values and full-vocabulary ranks to report. Strings (single tokens) or ints, or a sequence of them, are shared across every row. A tensor is per row: it is right-aligned and broadcast to the vectors’ leading dims (for [batch, pos, d_model] vectors pass [batch, pos] next-token ids, or [batch, 1] for one id per example) and the trailing axis of values has length 1. Note the convention differs from ActivationCache.logit_attrs, which reads a 1-D tensor as per-example first; here alignment is purely right-to-left.
return_type – "logits", "log_probs" or "probs".
top_k – Keep the k largest-logit tokens per vector, with their ids.
chunk_size – Rows of vectors processed per unembedding matmul.
enable_grad – Keep autograd on (for trainable lenses and attribution); default runs under torch.no_grad().

Returns:

A LogitReadout.

Raises:

ValueError – On a bad return_type, more than one selector, a chunk_size below 1, a wrong trailing dimension, empty vectors, per-row targets that do not broadcast, a model whose final norm cannot be applied exactly, or a bridge whose readout is not ln_final -> unembed on the last block’s output.
NotImplementedError – For encoder-decoder and encoder-only bridges.

transformer_lens.tools.analysis.mean_difference_basis(activations: Float[Tensor, 'batch d_model'], labels: Sequence[Hashable]) → Float[Tensor, 'd_model rank']¶

The span of the centred class means: for two classes, the difference-in-means direction.

Parameters:

activations – Activations, one row per example.
labels – The variable’s value for each example; at least two distinct values.

Returns:

Orthonormal basis of rank n_classes - 1, [d_model, n_classes - 1].

Raises:

ValueError – If fewer than two classes are present.

transformer_lens.tools.analysis.neuron_basis(d_mlp: int, neurons: Sequence[int] | ndarray | Tensor) → Float[Tensor, 'd_mlp k']¶

The coordinate basis of a set of MLP neurons, e.g. the features a sparse probe selects.

Interchanging it sets exactly those neurons to the source’s values, at a neuron hook such as blocks.N.mlp.hook_post. Compare it with random_neuron_basis() of the same size: a random direction is the wrong placebo for a claim about neurons, because it spreads over every neuron.

Parameters:

d_mlp – Number of neurons at the hook.
neurons – Integer indices of the neurons, each in [0, d_mlp) and none repeated.

Returns:

Orthonormal basis of unit vectors, [d_mlp, k], in the order given.

Raises:

ValueError – If no neuron is given, or an index is not an integer, out of range or repeated.

transformer_lens.tools.analysis.orthonormal_subspace(matrix: Float[Tensor, 'ambient width'], *, rank: int | None = None, rtol: float | None = None) → SubspaceBasis¶

Extract an explicitly ranked orthonormal column-space basis.

Low-precision inputs are promoted to float32 before the reduced SVD. With no explicit rtol, numerical rank uses the compute-SVD error scale relative to the largest singular value, so the default does not depend on the storage dtype — e.g. a 4096-row half-precision input gets 4096 * float32_eps (about 4.9e-4), tight enough to flag an exactly repeated column as rank-deficient. An explicit rank truncates the measured subspace but may not exceed its measured rank.

Parameters:

matrix – Finite floating-point matrix with shape [ambient_dim, width].
rank – Optional number of leading singular directions to retain.
rtol – Optional non-negative relative singular-value threshold.

Returns:

Basis, complete singular spectrum, and rank metadata.

Raises:

ValueError – If the matrix or rank policy is invalid.

transformer_lens.tools.analysis.probe_basis(weights: Float[Tensor, 'n_rows d_model']) → Float[Tensor, 'd_model rank']¶

The directions of a linear probe’s weights, in activation space.

Undo any input scaling first: for a probe fit on standardised inputs, divide each weight by its feature’s standard deviation.

Parameters:: weights – Probe weights; one row for a binary probe, one row per class otherwise.
Returns:: the weight direction for one row, the span of the centred rows (rank n_rows - 1) for several.
Return type:: Orthonormal basis

transformer_lens.tools.analysis.projection_kernel(subspace_a: SubspaceBasis, subspace_b: SubspaceBasis, *, check_orthonormal: bool = True) → ProjectionKernelResult¶

Measure overlap between two explicitly extracted subspaces.

Raw PK lies in [0, min(rank_a, rank_b)]. The normalized value is PK / sqrt(rank_a * rank_b), the cosine between the two projection matrices. Principal angles are returned in radians.

transformer_lens.tools.analysis.random_basis(d_model: int, rank: int, seed: int) → Float[Tensor, 'd_model rank']¶

A seeded random orthonormal basis: the baseline any aligned basis should beat.

Parameters:: seed – Seed for the Gaussian matrix the basis is drawn from.
Returns:: Orthonormal basis, [d_model, rank].

transformer_lens.tools.analysis.random_init_bridge(model_name: str, *, seed: int, dtype: dtype = torch.float32, device: str | device | None = None, **boot_kwargs: Any) → Any¶

A TransformerBridge with the architecture’s HuggingFace initialisation and no training.

The random-init control: train fit_das() on it with the same pairs, rank and budget as on the trained model. Effects DAS finds here come from the method and the architecture, not from anything the model learned. Built with load_weights=False, so the model class and prepared configuration match the trained bridge.

Parameters:

model_name – Model name or alias, as for TransformerBridge.boot_transformers.
device – Device for the bridge; by default the device boot_transformers would choose.
**boot_kwargs – Passed on to TransformerBridge.boot_transformers, e.g. revision.

Returns:

The untrained TransformerBridge, in evaluation mode.

transformer_lens.tools.analysis.random_neuron_basis(d_mlp: int, k: int, seed: int, exclude: Sequence[int] | ndarray | Tensor | None = None) → Float[Tensor, 'd_mlp k']¶

The coordinate basis of k neurons drawn at random: the placebo for neuron_basis().

Parameters:

d_mlp – Number of neurons at the hook.
seed – Seed for the draw; the caller’s random stream is untouched.
exclude – Neurons the draw must avoid, usually the ones the placebo is compared with.

Returns:

Orthonormal basis of unit vectors, [d_mlp, k].

Raises:

ValueError – If fewer than k neurons remain to draw from, or exclude holds an invalid index.

transformer_lens.tools.analysis.random_projection_kernel_moments(ambient_dim: int, rank: int) → RandomSubspaceReference¶

Return PK moments for independent Haar-distributed rank-rank planes.

These idealized descriptive moments are not calibrated p-values for trained model weights, whose head subspaces are dependent and anisotropic.

transformer_lens.tools.analysis.reinitialized(model: Module, seed: int) → Module¶

A copy of model with freshly initialised parameters: the random-init control.

For models built from standard nn layers, such as a HookedRootModule: every module with reset_parameters is reset, so every parameter must belong to one. Buffers, such as fixed architectural matrices, are kept. For HuggingFace models use random_init_bridge(), which applies the architecture’s own initialisation scheme.

Parameters:: model – The trained model; it is not modified.
Returns:: The untrained copy, on the device of the model’s parameters, in evaluation mode.
Raises:: ValueError – If a parameter belongs to no module with reset_parameters, so it would keep its trained value.

transformer_lens.tools.analysis.run_with_intervention(model: Any, hook_name: str, intervention: Callable[[Tensor, Tensor], Tensor], base_inputs: Tensor, source_activations: Float[Tensor, 'batch d_model'], task_logits: Callable[[Tensor], Float[Tensor, 'batch n_classes']], *, positions: Int[Tensor, 'batch'] | None = None) → Float[Tensor, 'batch n_classes']¶

Task logits for base_inputs with intervention applied at hook_name, in one pass.

Differentiable with respect to any parameters of intervention, for training an intervention. To evaluate one over many pairs in bounded memory, use intervened_logits().

Parameters:

model – A model exposing the hook API.
intervention – Maps base and source activations to new base activations, e.g. a SubspaceInterchange.
source_activations – Activations from the source inputs at the same hook point, one row per base, e.g. from collect_activations().
task_logits – Runs the model on a batch of inputs.
positions – For [batch, pos, d_model] activations, the position to intervene on per base.

Returns:

Task logits under the intervention.

Raises:

ValueError – If the activation does not match the positions given (or not given).
RuntimeError – If the hook does not fire exactly once in the forward pass.

transformer_lens.tools.analysis.solve_coordinate_patch(x: Tensor, dictionary: Tensor, source_idx: int, target_idx: int, *, decomposition: JSpaceDecomposition | None = None, k: int = 25, mode: str = 'substitute', alpha: float = 1.0, algorithm: str = 'nonnegative_orthogonal_matching_pursuit') → CoordinatePatch¶

Edit one sparse J-space coordinate while preserving its recovered frame.

substitute zeros the source and replaces the target coordinate with the source value; swap exchanges the two values. An absent target is appended at zero. The edit is blended by alpha and reconstructed over the original residual x - reconstruction.

Parameters:

x – Activation vector with shape [d_model].
dictionary – Atom matrix with shape [num_atoms, d_model]; rows are atoms.
source_idx – Atom index that must occur in the decomposition’s active support.
target_idx – Distinct atom index to receive or exchange the source coordinate.
decomposition – Optional decomposition of x under dictionary. Reuse validates only the selected and edited rows, avoiding another full-dictionary scan.
k – Sparse-solver upper bound when no decomposition is supplied.
mode – "substitute" to overwrite the target, or "swap" to exchange coordinates.
alpha – Finite interpolation strength. Zero returns an unchanged float32 clone.
algorithm – Sparse coefficient-update rule when no decomposition is supplied. A supplied decomposition is accepted if it fits as either algorithm’s output.

Returns:

A CoordinatePatch with the anchored edit and numerical diagnostics.

Raises:

ValueError – If tensors, indices, edit arguments, atoms, or a supplied decomposition are invalid or incompatible.
RuntimeError – If a fresh default NNLS decomposition fails its KKT certification, or an edit postcondition is violated.

Warns:

UserWarning – If the source/target pair is near-parallel, the normalized edit basis is rank deficient or poorly conditioned, or alpha extrapolates coordinates_after outside the nonnegative pursuit frame. All three are non-fatal.

transformer_lens.tools.analysis.sweep_sparse_probe(features: Float[Tensor, 'example feature'], labels: Bool[Tensor, 'example'] | Integer[Tensor, 'example'], *, ks: Sequence[int], test_fraction: int | float = 0.3, positive_label: int | bool = 1, groups: Integer[Tensor, 'example'] | None = None, preprocess: str = 'none', std_floor: int | float = 0.001, class_weight: str | None = 'balanced', l2_strength: int | float = 0.01, n_random_subsets: int = 0, n_label_shuffles: int = 0, seed: int = 0, max_iter: int = 200, max_refinement_steps: int = 25, decrement_tolerance: int | float = 1e-12) → SparseProbeSweep¶

Fit sparse probes and optional controls over one fixed train/test split.

ks must contain strictly increasing positive integers. Random-coordinate controls sample supports without replacement. Label-shuffle controls permute training labels, repeat selection and fitting, then evaluate against the untouched held-out labels. Control arrays contain raw metrics and do not represent automatic significance tests.

Parameters:

features – Finite float16/bfloat16/float32/float64 tensor shaped [example, feature].
labels – Boolean or integer binary labels shaped [example].
ks – Strictly increasing unique sparsity levels.
test_fraction – Requested held-out fraction within each class. With groups it targets held-out rows per class, which whole groups only approximate.
positive_label – Label defining the positive class and score sign.
groups – Optional integer group IDs; each group is held out whole, with both classes on both sides. Groups may mix classes, but each class must appear in two or more.
preprocess – "none" or train-only "standardize".
std_floor – Lower bound on a non-constant column’s "standardize" scale, shared by every fit; 0 disables it.
class_weight – "balanced" or None, shared by every fit.
l2_strength – Positive coefficient penalty shared by every fit.
n_random_subsets – Random-coordinate control fits per sparsity level.
n_label_shuffles – Shuffled-training-label control fits per sparsity level.
seed – Local CPU-generator seed for splitting and controls.
max_iter – Maximum LBFGS iterations per fit before Newton refinement. LBFGS also stops once its function evaluations reach max_iter * 5 // 4 (stop_reason="max_eval"; the line search can finish one past it), a budget equal to max_iter when max_iter <= 3.
max_refinement_steps – Maximum damped Newton steps per fit; zero only checks.
decrement_tolerance – Largest accepted Newton decrement g^T H^-1 g / 2 per fit, an estimate of the objective gap to the optimum in nats; a larger gap raises.

Returns:

Main probe results plus aligned raw control distributions. A control fit that fails to converge is excluded and recorded in rejections rather than aborting the sweep, so completed fits survive; each control’s rows cover only its converged repeats.

Raises:

ValueError – If the grid, controls, inputs, groups, or options are invalid.
SparseProbeConvergenceError – If a main-probe fit fails to converge. Control-fit convergence failures do not raise here — they are collected in SparseProbeSweep.rejections; any other error still propagates.

transformer_lens.tools.analysis package¶

Submodules¶

Module contents¶

Resolution order¶

Fit checkpoint schema (reference write_checkpoint() format)¶

Fit checkpoint schema (reference `write_checkpoint()` format)¶