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.

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). Integrated gradients (EAP-IG) and ablate-outside faithfulness are not implemented yet.

  • backward_lens: GPT-2 MLP weight-gradient factors projected into vocabulary space with explicit raw-gradient sign semantics.

  • direct_logit_attribution: Direct Logit Attribution (DLA) over components, layers, or attention heads.

  • direct_path_patching: Direct path patching for head-to-head circuit analysis.

  • 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).

  • jacobian_lens_causal_swap_benchmark: A causal coordinate-swap benchmark for coordinate_patch_hooks, with baseline-capability filtering, displacement-matched random-atom controls, exact Clopper-Pearson confidence intervals, and a versioned, fingerprinted JSON artifact schema. Its _cli sibling is the artifact-generation entry point and is deliberately not imported here, so python -m runs it once.

  • 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.

  • svd_circuits: Per-head QK/OV singular-vector decomposition with a degeneracy guard, OV vocab and logit readout, per-position activation projection onto singular directions, and a mandatory causal patch gate that reconstructs a head’s output onto a chosen singular subspace.

class transformer_lens.tools.analysis.ActivationProjection(head_svd: HeadSVD, coefficients: Float[Tensor, 'pos rank'], str_tokens: List[str])

Bases: object

Per-position coefficients of a head’s actual output in its OV output basis.

head_svd

The OV decomposition this was projected against.

Type:

transformer_lens.tools.analysis.svd_circuits.HeadSVD

coefficients

[pos, rank]; coefficients[:, i] is the signed amount of singular direction i (head_svd.V[:, i]) present in the head’s actual output at each position. Summing coefficients[:, i] * head_svd.V[:, i] over i reconstructs the head’s real per-position output to numerical precision, since V’s columns are orthonormal and this projects onto the exact basis the head writes in.

Type:

jaxtyping.Float[Tensor, ‘pos rank’]

str_tokens

Tokenized prompt, aligned with the position axis, for display.

Type:

List[str]

coefficients: Float[Tensor, 'pos rank']
head_svd: HeadSVD
str_tokens: List[str]
class transformer_lens.tools.analysis.AttentionHeadRef(layer: int, head: int, role: Literal['Q', 'K', 'V', 'O'], kind: Literal['query', 'kv'])

Bases: object

Structured identity for one attention-head weight subspace.

head: int
kind: Literal['query', 'kv']
property label: str

Return the conventional TransformerLens layer/head label.

layer: int
role: Literal['Q', 'K', 'V', 'O']
class transformer_lens.tools.analysis.AttributionResult(node_scores: dict[~transformer_lens.tools.analysis.attribution_patching.Node, float], edge_scores: dict[tuple[~transformer_lens.tools.analysis.attribution_patching.Node, ~transformer_lens.tools.analysis.attribution_patching.Node], float] = <factory>)

Bases: object

Scored output of an attribution-patching sweep.

node_scores

Signed first-order effect estimate per node, (a_clean - a_corrupt) . d(metric)/d(a). A positive score means patching that node from corrupt toward clean moves the metric in the positive direction (the denoising convention pinned in the module docstring). For an edge-granularity sweep this is instead each writer’s aggregate over its own outgoing edge scores. Because the graph has a terminal logits reader that every writer feeds, that aggregate equals the writer’s direct node score (see attribution_patch()).

Type:

dict[transformer_lens.tools.analysis.attribution_patching.Node, float]

edge_scores

Per-edge effect estimate keyed by (source, destination). Populated for an edge-granularity sweep; empty for a node sweep.

Type:

dict[tuple[transformer_lens.tools.analysis.attribution_patching.Node, transformer_lens.tools.analysis.attribution_patching.Node], float]

edge_scores: dict[tuple[Node, Node], float]
node_scores: dict[Node, float]
top_edges(k: int = 10) → list[tuple[Node, Node, float]]

The k edges with the largest effect magnitude, strongest first.

Ranking is by absolute score, matching top_nodes: a large negative edge effect is as causally important as a large positive one. Ties keep enumeration order (stable sort). Requesting more than the available edges returns all of them.

top_nodes(k: int = 10) → list[tuple[Node, float]]

The k nodes with the largest effect magnitude, strongest first.

Ranking is by absolute score: a node with a large negative effect is as causally important as one with a large positive effect, so magnitude — not signed value — orders the circuit. Ties keep enumeration order (stable sort). Requesting more than the available nodes returns all of them.

class transformer_lens.tools.analysis.BackwardLens(model: Any)

Bases: object

Analyze dense MLP weight gradients in the output vocabulary basis.

The analyzer accepts a fresh, raw dense-MLP TransformerBridge such as GPT-2 or Pythia/GPT-NeoX. Results retain no model or tokenizer reference and contain detached CPU-owned tensors. Raw backward signals are loss gradients; gradient descent subtracts them.

__init__(model: Any)

Validate and retain the raw dense-MLP Bridge used for analyses.

analyze(prompt: str, target_token: str, layers: Sequence[int], *, normalized: bool = False, top_k: int = 10, return_full_logits: bool = False) → BackwardLensResult

Analyze one final-position, one-token target loss.

Parameters:

  • prompt – Non-empty unbatched prompt text.

  • target_token – Text encoding to exactly one token without BOS.

  • layers – Unique layer indices in desired result order.

  • normalized – Also project unit-normalized nonzero factors using the Normalized Logit Lens. Raw projections are always returned.

  • top_k – Number of largest and smallest values and token ids retained per matrix and position. Defaults to 10.

  • return_full_logits – Also retain full vocabulary tensors on CPU. Defaults to False to keep result size bounded.

Returns:

Detached gradient factors, bounded vocabulary rankings, norms, reconstruction errors, target metadata, and optional full logits.

class transformer_lens.tools.analysis.BackwardLensLayerResult(layer: int, input_projection: BackwardLensMatrixResult, output_projection: BackwardLensMatrixResult)

Bases: object

Vocabulary-facing input/output MLP matrix results for one indexed layer.

input_projection: BackwardLensMatrixResult
layer: int
output_projection: BackwardLensMatrixResult
class transformer_lens.tools.analysis.BackwardLensMatrixResult(factors: LinearGradientFactors, projected_factor: Literal['forward_inputs', 'output_gradients'], factor_norms: Float[Tensor, 'position'], zero_norm_mask: Bool[Tensor, 'position'], vocabulary_size: int, target_token_id: int, top_ranking: VocabularyRanking, bottom_ranking: VocabularyRanking, target_largest_ranks: Int[Tensor, 'position'], target_smallest_ranks: Int[Tensor, 'position'], normalized_top_ranking: VocabularyRanking | None = None, normalized_bottom_ranking: VocabularyRanking | None = None, normalized_target_largest_ranks: Int[Tensor, 'position'] | None = None, normalized_target_smallest_ranks: Int[Tensor, 'position'] | None = None, vocabulary_logits: Float[Tensor, 'position d_vocab'] | None = None, normalized_vocabulary_logits: Float[Tensor, 'position d_vocab'] | None = None)

Bases: object

Factors and vocabulary readouts for one dense MLP weight matrix.

factors contains the full linear factorization. projected_factor says whether its residual-width forward_inputs or raw-gradient output_gradients were decoded. factor_norms and zero_norm_mask have shape [position] with float32 and bool dtypes. Largest and smallest signed rankings are always retained. Full vocabulary_logits are present only when explicitly requested. Normalized rankings and optional full logits are present when the Normalized Logit Lens is requested. Every retained tensor is a detached CPU-owned value; gradient descent subtracts raw gradients.

bottom(*, k: int, normalized: bool = False) → VocabularyRanking

Return up to the retained smallest signed logits and token ids.

bottom_ranking: VocabularyRanking
bottom_tokens(tokenizer: Any, *, k: int, normalized: bool = False) → list[list[str]]

Decode the smallest-k vocabulary ids for every position.

factor_norms: Float[Tensor, 'position']
factors: LinearGradientFactors
gradient_descent_target_ranks(target_token_id: int, *, normalized: bool = False) → Int[Tensor, 'position']

Return ascending raw-gradient target ranks (rank zero is smallest).

logits(*, normalized: bool = False) → Float[Tensor, 'position d_vocab']

Return opted-in raw or Normalized Logit Lens full logits.

normalized_bottom_ranking: VocabularyRanking | None = None
normalized_target_largest_ranks: Int[Tensor, 'position'] | None = None
normalized_target_smallest_ranks: Int[Tensor, 'position'] | None = None
normalized_top_ranking: VocabularyRanking | None = None
normalized_vocabulary_logits: Float[Tensor, 'position d_vocab'] | None = None
projected_factor: Literal['forward_inputs', 'output_gradients']
target_largest_ranks: Int[Tensor, 'position']
target_ranks(target_token_id: int, *, largest: bool, normalized: bool = False) → Int[Tensor, 'position']

Return zero-based target ranks per position in the requested ordering.

largest=True gives rank zero to the largest logit. largest=False gives rank zero to the smallest, which is the useful raw-gradient convention for the second MLP matrix because gradient descent subtracts it. Ties receive the same competition rank. The analyzed target’s ranks are always retained; other token ids require opted-in full logits.

target_smallest_ranks: Int[Tensor, 'position']
target_token_id: int
top(*, k: int, normalized: bool = False) → VocabularyRanking

Return up to the retained largest signed logits and token ids.

top_ranking: VocabularyRanking
top_tokens(tokenizer: Any, *, k: int, normalized: bool = False) → list[list[str]]

Decode the largest-k vocabulary ids for every position.

vocabulary_logits: Float[Tensor, 'position d_vocab'] | None = None
vocabulary_size: int
zero_norm_mask: Bool[Tensor, 'position']
class transformer_lens.tools.analysis.BackwardLensResult(prompt: str, prompt_token_ids: Int[Tensor, 'position'], target_token: str, target_token_id: int, loss: float, layers: tuple[BackwardLensLayerResult, ...], max_absolute_reconstruction_error: float, max_relative_reconstruction_error: float, includes_normalized_logits: bool, includes_full_logits: bool)

Bases: object

Detached result of one BackwardLens.analyze() call.

prompt and target_token echo the analyzed inputs; target_token_id is the single vocabulary id the target text encodes to. loss is the raw scalar cross-entropy of the final-position next-token prediction against the target; it preserves the d(loss)/d(...) sign convention and is not negated. prompt_token_ids is an owned CPU int64 tensor with shape [position]; every residual-width factor in layers is aligned to these same positions. Position zero is a prepended BOS only when the model and tokenizer configuration requests one. layers preserves requested order. Maximum errors summarize both matrices over every requested layer. includes_normalized_logits records whether the Normalized Logit Lens was computed. includes_full_logits records whether full vocabulary tensors were retained in addition to bounded rankings. No model or tokenizer reference is retained.

includes_full_logits: bool
includes_normalized_logits: bool
layer(layer: int) → BackwardLensLayerResult

Return one requested layer result or raise KeyError.

layers: tuple[BackwardLensLayerResult, ...]
loss: float
max_absolute_reconstruction_error: float
max_relative_reconstruction_error: float
prompt: str
prompt_token_ids: Int[Tensor, 'position']
target_token: str
target_token_id: int
class transformer_lens.tools.analysis.BenchmarkCorpus(name: str, concepts: Sequence[str], functions: Sequence[FunctionSpec])

Bases: object

A named set of concepts and the prompt functions evaluated over them.

concepts: Sequence[str]
functions: Sequence[FunctionSpec]
name: str
class transformer_lens.tools.analysis.CoordinatePatch(support_before: Tensor, support_after: Tensor, coordinates_before: Tensor, coordinates_after: Tensor, source_slot: int, target_slot: int, target_was_appended: bool, target_was_selected: bool, overwritten_target_coordinate: float | None, reconstruction_before: Tensor, reconstruction_after: Tensor, residual: Tensor, delta: Tensor, patched: Tensor, nonedited_coordinate_max_delta: float, residual_max_delta: float, source_target_cosine: float, basis_rank: int, basis_condition_number: float, coordinates_after_nonnegative: bool)

Bases: object

Result of an anchored edit to sparse J-space coordinates.

All tensors are detached. Vector outputs use float32 on the dictionary device; the two support tensors are torch.long on CPU. This is a report only: it retains no model, lens, tokenizer, hook, or full dictionary.

support_before

Active support from the sparse decomposition, in decomposition order.

Type:

torch.Tensor

support_after

The coordinate frame the edit acts on: support_before in the same order, with an absent target appended. Both coordinate tensors align with this frame.

Type:

torch.Tensor

coordinates_before

Original coordinates in support_after, including an appended zero when the target was absent.

Type:

torch.Tensor

coordinates_after

Coordinates after applying mode and alpha.

Type:

torch.Tensor

source_slot

Position of the source atom within support_after.

Type:

int

target_slot

Position of the target atom within support_after.

Type:

int

target_was_appended

Whether the target was absent from the active support and appended to support_after. target_was_selected disambiguates why: pursuit may never have considered the target at all, or may have selected it and then assigned it a zero coordinate.

Type:

bool

target_was_selected

Whether the target atom was selected by pursuit (decomposition.selected_support), whether or not it ended up numerically active. target_was_appended and not target_was_selected means pursuit never considered the target; target_was_appended and target_was_selected means pursuit selected it but assigned it a zero coordinate. Always True when target_was_appended is False, since the active support is a subset of the selected support.

Type:

bool

overwritten_target_coordinate

The old target coordinate that a substitute discards when it overwrites an already-active target; None when the target was absent or mode is "swap" (a swap relocates a coordinate and discards nothing). Alpha-independent by design, like target_was_appended and the support fields: it reports the coordinate the requested edit targets for replacement, not the blended outcome, so it is non-None even at alpha=0 where nothing is actually discarded.

Type:

float | None

reconstruction_before

Sparse reconstruction from support_before.

Type:

torch.Tensor

reconstruction_after

reconstruction_before + delta, the anchored reconstruction after the edit.

Type:

torch.Tensor

residual

Anchored residual x - reconstruction_before. This is not necessarily the decomposition’s orthogonal non_j_space_component.

Type:

torch.Tensor

delta

basis @ (coordinates_after - coordinates_before); equivalently reconstruction_after - reconstruction_before. Exactly zero and proportional to alpha, so it carries no recompute floor.

Type:

torch.Tensor

patched

x + delta; equal to the float32 input when alpha is zero, where delta is then exactly zero.

Type:

torch.Tensor

nonedited_coordinate_max_delta

Maximum absolute change across every coordinate that is neither source nor target. A postcondition witness: it must be numerically zero.

Type:

float

residual_max_delta

Maximum absolute difference between residual and patched - dictionary[support_after].T @ coordinates_after, the reconstruction recomputed independently from the dictionary rows. A postcondition witness: the anchored residual must be preserved.

Type:

float

source_target_cosine

Signed cosine between the source and target atoms.

Type:

float

basis_rank

Numerical rank of the column-normalized edit basis at float32 precision.

Type:

int

basis_condition_number

Condition number of that basis, or infinity when rank deficient.

Type:

float

coordinates_after_nonnegative

Whether every entry of coordinates_after is >= 0. alpha in [0, 1] interpolates within the nonnegative pursuit frame and always leaves this True; an alpha outside that range (e.g. 2.0 or -1.0) extrapolates and can drive a coordinate negative, in which case this is False and a warning is raised.

Type:

bool

basis_condition_number: float
basis_rank: int
coordinates_after: Tensor
coordinates_after_nonnegative: bool
coordinates_before: Tensor
delta: Tensor
nonedited_coordinate_max_delta: float
overwritten_target_coordinate: float | None
patched: Tensor
reconstruction_after: Tensor
reconstruction_before: Tensor
residual: Tensor
residual_max_delta: float
source_slot: int
source_target_cosine: float
support_after: Tensor
support_before: Tensor
target_slot: int
target_was_appended: bool
target_was_selected: bool
exception transformer_lens.tools.analysis.DegenerateDirectionError

Bases: ValueError

Raised when per-direction attribution is requested for a degenerate direction.

Subclasses ValueError so callers that already except ValueError keep working, mirroring how the other analysis tools raise ValueError for their input guards.

class transformer_lens.tools.analysis.DirectLogitAttribution(attribution: Float[Tensor, 'component *batch_and_pos'], labels: List[str], unit: str)

Bases: object

Result of a direct_logit_attribution() call.

attribution

Tensor of logit (or logit-difference) attributions with shape [component, *batch_and_pos]. The leading axis is aligned with labels. When pos selects a single position (the default) the position axis is dropped, leaving [component, batch] — or [component] if the cache had its batch dimension removed.

Type:

jaxtyping.Float[Tensor, ‘component *batch_and_pos’]

labels

Human-readable name for each component, aligned with the leading axis of attribution (e.g. "embed", "0_attn_out", "L3H7").

Type:

List[str]

unit

The decomposition unit used (“component”, “layer”, or “head”).

Type:

str

attribution: Float[Tensor, 'component *batch_and_pos']
labels: List[str]
top(k: int = 5) → List[tuple]

Return the k highest-attribution (label, value) pairs.

Attribution is reduced to a scalar per component by meaning over any remaining batch/position dimensions, so this is most meaningful when a single position was selected.

unit: str
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 yet. 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.FunctionSpec(name: str, template: str, answers: Dict[str, str])

Bases: object

A templated prompt function evaluated over a shared set of concepts.

template takes a single {arg} placeholder, e.g. "The capital of {arg} is". answers maps each concept to its answer word under this function, e.g. {"France": "Paris"}.

answers: Dict[str, str]
name: str
template: str
class transformer_lens.tools.analysis.HeadAffinityPair(source: AttentionHeadRef, target: AttentionHeadRef, score: float, normalized: float)

Bases: object

One ranked source-target head pair.

normalized: float
score: float
source: AttentionHeadRef
target: AttentionHeadRef
class transformer_lens.tools.analysis.HeadAffinityResult(scores: Float[Tensor, 'source_layer source_head target_layer target_head'], normalized: Float[Tensor, 'source_layer source_head target_layer target_head'], valid_mask: Bool[Tensor, 'source_layer source_head target_layer target_head'], source_role: Literal['Q', 'K', 'V', 'O'], target_role: Literal['Q', 'K', 'V', 'O'], source_layer_indices: Tuple[int, ...], target_layer_indices: Tuple[int, ...], source_head_kind: Literal['query', 'kv'], target_head_kind: Literal['query', 'kv'], source_ranks: Int[Tensor, 'source_layer source_head'], target_ranks: Int[Tensor, 'target_layer target_head'], source_rank: int, target_rank: int, rank: int | None, rtol: float)

Bases: object

Projection Kernel affinities between two attention-head roles.

Score tensors have shape [source_layer, source_head, target_layer, target_head]. Layer index tuples map tensor positions to original model block numbers.

source_ranks and target_ranks are measured numerical ranks for each head before optional truncation. Scalar source_rank and target_rank are the retained basis widths used for their respective roles.

normalized: Float[Tensor, 'source_layer source_head target_layer target_head']
rank: int | None
rtol: float
scores: Float[Tensor, 'source_layer source_head target_layer target_head']
source_head_kind: Literal['query', 'kv']
source_layer_indices: Tuple[int, ...]
source_rank: int
source_ranks: Int[Tensor, 'source_layer source_head']
source_role: Literal['Q', 'K', 'V', 'O']
target_head_kind: Literal['query', 'kv']
target_layer_indices: Tuple[int, ...]
target_rank: int
target_ranks: Int[Tensor, 'target_layer target_head']
target_role: Literal['Q', 'K', 'V', 'O']
top_pairs(k: int = 20, *, normalized: bool = False) → List[HeadAffinityPair]

Return the highest-scoring valid pairs with deterministic tie order.

valid_mask: Bool[Tensor, 'source_layer source_head target_layer target_head']
class transformer_lens.tools.analysis.HeadDecomposition(layer: int, head: int, QK: HeadSVD | None = None, OV: HeadSVD | None = None)

Bases: object

Container returned by decompose_head().

QK and OV hold the HeadSVD for each requested map, or None when that map was not requested.

OV: HeadSVD | None = None
QK: HeadSVD | None = None
head: int
layer: int
class transformer_lens.tools.analysis.HeadSVD(which: Literal['QK', 'OV'], layer: int, head: int, U: Float[Tensor, 'd_model rank'], S: Float[Tensor, 'rank'], V: Float[Tensor, 'd_model rank'], rank_report: List[RankReportRow], eps: float, null_rtol: float, folded_ln: bool = False)

Bases: object

Factored SVD of one head map ("QK" or "OV") with a degeneracy report.

which

Which map this decomposes, "QK" or "OV".

Type:

Literal[‘QK’, ‘OV’]

layer

Layer of the decomposed head.

Type:

int

head

Head index within the layer.

Type:

int

U

Left singular vectors, [d_model, rank]: column i is the map’s input direction i (for OV, the residual-stream direction this head’s value computation reads from; for QK, the destination/query-read direction).

Type:

jaxtyping.Float[Tensor, ‘d_model rank’]

S

Singular values, [rank], sorted descending.

Type:

jaxtyping.Float[Tensor, ‘rank’]

V

Right singular vectors, [d_model, rank]: column i is the map’s output direction i for OV (the residual-stream direction this head writes into, the one to project through W_U), or the source/key-read direction for QK (QK produces no write direction). The reconstruction is U @ S.diag() @ V.transpose(-2, -1).

Type:

jaxtyping.Float[Tensor, ‘d_model rank’]

rank_report

Per-direction RankReportRow list, aligned with the columns of U/V.

Type:

List[transformer_lens.tools.analysis.svd_circuits.RankReportRow]

eps

Relative gap below which adjacent directions share a block; every block boundary sits at a gap of at least eps.

Type:

float

null_rtol

Relative-to-top-singular-value tolerance below which a direction is numerically null.

Type:

float

folded_ln

Whether the model’s weights carried a folded final LayerNorm when this was decomposed. enable_compatibility_mode folds ln1 into W_V and centres W_O/W_U, so a decomposition describes the OV map only under the state it was built in; the readout and patch consumers refuse a decomposition whose state no longer matches the model.

Type:

bool

S: Float[Tensor, 'rank']
U: Float[Tensor, 'd_model rank']
V: Float[Tensor, 'd_model rank']
block_of(i: int) → List[int]

Return every direction index sharing direction i’s degeneracy block.

The result is a singleton [i] for an isolated direction and the full run for a degenerate one.

degenerate_blocks() → List[List[int]]

Return every degenerate block’s indices: the subspaces to attribute whole or skip.

eps: float
folded_ln: bool = False
head: int
is_degenerate(i: int) → bool

Whether direction i is refused: rotation-ambiguous inside a block, or null.

layer: int
null_rtol: float
rank_report: List[RankReportRow]
require_isolated(i: int) → None

Raise DegenerateDirectionError unless direction i is attributable alone.

The message names the cause, since a rotation-ambiguous block is still a subspace worth attributing while a null block carries no signal.

which: Literal['QK', 'OV']
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.

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.JSpaceOccupancy(occupancy: int, marginal_captured_variance: Tensor, control_captured_variance: Tensor, support: Tensor)

Bases: object

Result of a J-space occupancy estimate.

occupancy

Estimated number of meaningfully-active atoms – the step of maximum separation between the real and random-control cumulative captured variance.

Type:

int

marginal_captured_variance

Per-step captured-variance gain of the real greedy selection, shape [max_atoms].

Type:

torch.Tensor

control_captured_variance

Per-step captured-variance gain averaged over the random control dictionaries, shape [max_atoms].

Type:

torch.Tensor

support

Greedily selected atom indices, shape [max_atoms] (token ids when the dictionary is the vocabulary of J-lens vectors).

Type:

torch.Tensor

control_captured_variance: Tensor
marginal_captured_variance: Tensor
occupancy: int
support: Tensor
class transformer_lens.tools.analysis.JSpaceVarianceProfile(layers: List[int], median: Dict[int, float], pooled: Dict[int, float], per_position: Dict[int, Tensor])

Bases: object

Per-layer J-space variance profile over a prompt corpus.

Produced by JacobianLens.fraction_of_variance().

layers

The source layers profiled, in order.

Type:

List[int]

median

Per-layer median over positions of the J-space variance fraction ||j_space_component||^2 / ||activation||^2.

Type:

Dict[int, float]

pooled

Per-layer pooled ratio sum(||j_space_component||^2) / sum(||activation||^2) across the corpus (the paper’s “fraction of total variance”).

Type:

Dict[int, float]

per_position

Per-layer 1-D tensor of the raw per-position variance fractions.

Type:

Dict[int, torch.Tensor]

layers: List[int]
median: Dict[int, float]
per_position: Dict[int, Tensor]
pooled: Dict[int, float]
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.

For each token’s unit lens vector v̂: h <- h - (h·v̂) v̂, applied sequentially when several tokens are given.

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 solving a fresh decomposition.

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.

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, 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, or a token tensor that does not have shape [1, seq].

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

  1. Local file — if name_or_path is an existing .pt file, load it directly.

  2. Local directory — if name_or_path is a directory, load <name_or_path>/<filename>.

  3. 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. The filename argument is ignored in this case because the registry already encodes the correct subpath.

  4. 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) to pin. When omitted, the Hub repository’s mutable default branch is followed; pin a commit hash for reproducible analyses.

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) → JacobianLens

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 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 or mismatched lenses.

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().

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, or a requested layer has no transport matrix.

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]]

Hooks that clamp lens coordinates to their clean-run exchange.

For each layer, this projects the corresponding activation from clean_cache into that layer’s lens basis, exchanges its source and target coordinates once, and holds the live activation at that fixed target. Unlike swap_hooks(), 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.

  • clean_cache – 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.

  • positions – Chunk-local positions to clamp (negative indices allowed and normalized against the clean activations). Defaults to all.

Returns:

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

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

Hooks that swap two concepts’ live 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 (sigma(c) - c) where sigma exchanges the two coordinates. The component of h orthogonal to span{v_s, v_t} is untouched. alpha=2 is the paper’s “double-strength” swap.

This transform re-reads c 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. For the paper’s multi-layer clamp protocol, use swap_clamp_hooks() with activations cached from the clean run.

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.

  • positions – Chunk-local positions to swap (negative indices allowed and normalized on every hook invocation). Defaults to all.

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.JacobianLensReadout(lens_topk_values: Dict[int, Float[Tensor, 'pos k']], lens_topk_indices: Dict[int, Int[Tensor, 'pos k']], model_topk_values: Float[Tensor, 'pos k'], model_topk_indices: Int[Tensor, 'pos k'], tokens: Int[Tensor, 'seq'], positions: List[int], use_jacobian: bool = True, lens_logits: Dict[int, Float[Tensor, 'pos d_vocab']] | None = None, model_logits: Float[Tensor, 'pos d_vocab'] | None = None)

Bases: object

Result of a JacobianLens.readout() call.

lens_topk_values

Per-layer retained top-k pre-softmax values, on CPU.

Type:

Dict[int, jaxtyping.Float[Tensor, ‘pos k’]]

lens_topk_indices

Per-layer retained top-k vocabulary ids, on CPU.

Type:

Dict[int, jaxtyping.Int[Tensor, ‘pos k’]]

model_topk_values

The model output’s retained top-k pre-softmax values, on CPU.

Type:

jaxtyping.Float[Tensor, ‘pos k’]

model_topk_indices

The model output’s retained top-k vocabulary ids, on CPU.

Type:

jaxtyping.Int[Tensor, ‘pos k’]

lens_logits

Optional full per-layer logits, on CPU. Present only when readout(return_full_logits=True) was requested.

Type:

Dict[int, jaxtyping.Float[Tensor, ‘pos d_vocab’]] | None

model_logits

Optional full model logits, on CPU. Present only when readout(return_full_logits=True) was requested.

Type:

jaxtyping.Float[Tensor, ‘pos d_vocab’] | None

tokens

The token ids of the run prompt, [seq].

Type:

jaxtyping.Int[Tensor, ‘seq’]

positions

The (normalized, non-negative) positions the readout covers, aligned with the pos axis of retained top-k and optional full logits.

Type:

List[int]

use_jacobian

Whether the Jacobian transport was applied (False = logit lens).

Type:

bool

lens_logits: Dict[int, Float[Tensor, 'pos d_vocab']] | None = None
lens_topk_indices: Dict[int, Int[Tensor, 'pos k']]
lens_topk_values: Dict[int, Float[Tensor, 'pos k']]
model_logits: Float[Tensor, 'pos d_vocab'] | None = None
model_topk_indices: Int[Tensor, 'pos k']
model_topk_values: Float[Tensor, 'pos k']
positions: List[int]
tokens: Int[Tensor, 'seq']
top_tokens(tokenizer: Any, k: int = 5) → Dict[int, List[List[str]]]

Decode the top-k tokens per layer and position.

Parameters:

  • tokenizer – The model’s tokenizer (model.tokenizer).

  • k – Number of top tokens to decode per (layer, position).

Returns:

{layer: [ [top-k strings] per position ]}, positions aligned with positions.

use_jacobian: bool = True
class transformer_lens.tools.analysis.LinearGradientFactors(forward_inputs: Float[Tensor, 'position in_features'], output_gradients: Float[Tensor, 'position out_features'], weight_gradient: Float[Tensor, 'weight_dim_0 weight_dim_1'], reconstructed_gradient: Float[Tensor, 'weight_dim_0 weight_dim_1'], absolute_reconstruction_error: float, relative_reconstruction_error: float, weight_layout: Literal['in_out', 'out_in'])

Bases: object

Detached factors and reconstruction for one linear weight gradient.

forward_inputs and output_gradients have shapes [position, in] and [position, out]. output_gradients and weight_gradient preserve the raw d(loss)/d(tensor) sign; they are not negated into update directions. Gradient tensors use the requested storage layout. All tensors are cloned to CPU in float32 so the result owns no autograd graph.

absolute_reconstruction_error: float
forward_inputs: Float[Tensor, 'position in_features']
output_gradients: Float[Tensor, 'position out_features']
reconstructed_gradient: Float[Tensor, 'weight_dim_0 weight_dim_1']
relative_reconstruction_error: float
weight_gradient: Float[Tensor, 'weight_dim_0 weight_dim_1']
weight_layout: Literal['in_out', 'out_in']
class transformer_lens.tools.analysis.LogitLensResult(readout: LogitReadout, labels: List[str], has_batch_dim: bool, target_names: List[str | None] | None = None)

Bases: object

Result of logit_lens().

readout

The LogitReadout over the accumulated stack. Its leading axis is the stack entry, aligned with labels; the remaining leading dims are [batch, pos] after the requested positions / batch_slice (an integer selection drops its axis).

Type:

transformer_lens.tools.analysis.logit_lens.LogitReadout

labels

Stack entry labels from accumulated_resid ("0_pre", "3_mid", "final_post").

Type:

List[str]

has_batch_dim

Whether a batch axis follows the stack axis.

Type:

bool

target_names

For shared targets, the string passed for each target (None where an id was passed), so rank_trajectory() accepts the strings again.

Type:

List[str | None] | None

has_batch_dim: bool
labels: List[str]
rank_trajectory(target: str | int | None = None) → Tensor

Full-vocabulary rank of one target at every stack entry, [layer, ...].

target is the string passed to targets or a token id; None selects the only target when exactly one was requested, and is the only form for per-row targets.

readout: LogitReadout
target_names: List[str | None] | None = None
to_dataframe(tokenizer: Any | None = None, top_k: int | None = None) → Any

One row per (stack entry, batch, position, token) as a pandas DataFrame.

Rows within a (layer, batch, pos) group are sorted by descending value and top_k keeps the leading ones. Full-vocabulary "logits" readouts go through transformer_lens.utilities.logits_utils.logits_to_df() (columns logit, log_prob, probability); every other readout carries token_index, optional token_string and one value column named like logits_to_df names it (logit / log_prob / probability).

top_tokens(tokenizer: Any) → Dict[str, Any]

Decoded top-k tokens per stack entry, {label: nested lists}.

property values: Tensor

Shortcut for readout.values.

class transformer_lens.tools.analysis.LogitReadout(values: Float[Tensor, '... selected'], return_type: str, entropy: Float[Tensor, '...'], top_ids: Int[Tensor, '... k'] | None = None, vocab_ids: Int[Tensor, 'v'] | None = None, target_ids: Tensor | None = None, target_ranks: Int[Tensor, '... t'] | None = None, per_row_targets: bool = False, applied_ln: bool = True, applied_output_transform: bool = True)

Bases: object

Result of logit_readout().

Exactly one of the three selectors expands the trailing axis of values: targets gives [..., n_targets] ([..., 1] for per-row targets), top_k gives [..., k] with top_ids, vocab gives [..., len(vocab)]; with none of them the trailing axis is the full vocabulary. entropy and target_ranks are always computed over the full vocabulary, whatever was selected.

values

The requested quantity (see return_type) with the leading dims of the input vectors and the selected trailing axis.

Type:

jaxtyping.Float[Tensor, ‘… selected’]

return_type

"logits", "log_probs" or "probs".

Type:

str

entropy

Full-vocabulary entropy of the softmax, in nats, [...].

Type:

jaxtyping.Float[Tensor, ‘…’]

top_ids

Token ids aligned with values when top_k was used.

Type:

jaxtyping.Int[Tensor, ‘… k’] | None

vocab_ids

The requested subset ids when vocab was used.

Type:

jaxtyping.Int[Tensor, ‘v’] | None

target_ids

Shared targets as [n_targets]; per-row targets as a tensor with the leading dims of the vectors (one id per row).

Type:

torch.Tensor | None

target_ranks

Zero-based full-vocabulary rank of each target (0 is the argmax; ties count only strictly larger logits), aligned with values.

Type:

jaxtyping.Int[Tensor, ‘… t’] | None

per_row_targets

Whether targets was given per row (a tensor).

Type:

bool

applied_ln

Whether the model’s final norm was applied to the vectors.

Type:

bool

applied_output_transform

Whether the adapter’s post-unembedding transform was invoked (it is the identity for most architectures).

Type:

bool

applied_ln: bool = True
applied_output_transform: bool = True
decode(tokenizer: Any) → Any

Decode the trailing-axis token ids to strings, keeping the leading structure.

top_k and per-row readouts decode per row; shared targets / vocab readouts decode the id list once. Full-vocabulary readouts have no id list to decode.

entropy: Float[Tensor, '...']
per_row_targets: bool = False
return_type: str
row_ids(index: Tuple[int, ...]) → Tensor | None

Trailing-axis token ids for one row of values (None = full vocabulary).

property selected_ids: Tensor | None

Token ids along the trailing axis of values when it is not the full vocabulary.

target_ids: Tensor | None = None
target_ranks: Int[Tensor, '... t'] | None = None
top_ids: Int[Tensor, '... k'] | None = None
values: Float[Tensor, '... selected']
vocab_ids: Int[Tensor, 'v'] | None = None
class transformer_lens.tools.analysis.LogitSignature(direction: int, values: Float[Tensor, 'token'])

Bases: object

Rank-1-reconstruction logit effect for one OV direction, per requested token.

direction

Which HeadSVD column this reconstructs.

Type:

int

values

Signed logit contribution, aligned with the requested tokens.

Type:

jaxtyping.Float[Tensor, ‘token’]

direction: int
values: Float[Tensor, 'token']
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.PatchResult(head_svd: HeadSVD, retained: List[int], original_metric: float, patched_metric: float, delta_metric: float, baseline_delta_metric: float, gated: bool)

Bases: object

Result of causally patching a head’s output onto a chosen OV singular subspace.

head_svd

The OV decomposition patched against.

Type:

transformer_lens.tools.analysis.svd_circuits.HeadSVD

retained

Direction indices whose span the head’s output was reconstructed onto; the complement was zeroed.

Type:

List[int]

original_metric

Metric value on the unmodified prompt.

Type:

float

patched_metric

Metric value after the subspace reconstruction.

Type:

float

delta_metric

patched_metric - original_metric.

Type:

float

baseline_delta_metric

the mean of the per-draw delta_metric magnitudes over several random control subspaces of the same width as retained, each drawn inside the head’s own OV span span(V) rather than from the full residual stream, so the control is the effect of an arbitrary same-size subspace of this head’s output rather than of an unrelated residual-stream direction. Averaging the magnitudes rather than the signed deltas keeps this a typical control effect that does not shrink when the controls mix sign, so it is non-negative.

Type:

float

gated

whether the retained subspace passed the causal test for the mode it was expressed in, comparing abs(delta_metric) against baseline_delta_metric (or an explicit threshold, if one was passed). For ablate (retain the complement), removing a load-bearing subspace should move the metric more than removing an arbitrary same-size one, so gated is abs(delta_metric) > threshold. For keep (retain only the given subspace), a subspace that reconstructs the head’s behavior should move the metric less than keeping an arbitrary same-size one, so gated is abs(delta_metric) < threshold. A single “moved more than baseline” test cannot answer both, since keep=S and ablate=complement(S) resolve to the same retained set.

Type:

bool

baseline_delta_metric: float
delta_metric: float
gated: bool
head_svd: HeadSVD
original_metric: float
patched_metric: float
retained: List[int]
class transformer_lens.tools.analysis.ProjectionKernelResult(score: Float[Tensor, ''], normalized: Float[Tensor, ''], cosines: Float[Tensor, 'principal_angle'], angles: Float[Tensor, 'principal_angle'], rank_a: int, rank_b: int, ambient_dim: int)

Bases: object

Projection Kernel score and its principal-angle decomposition.

ambient_dim: int
angles: Float[Tensor, 'principal_angle']
cosines: Float[Tensor, 'principal_angle']
normalized: Float[Tensor, '']
rank_a: int
rank_b: int
score: Float[Tensor, '']
class transformer_lens.tools.analysis.RandomSubspaceReference(ambient_dim: int, rank: int, mean: float, variance: float)

Bases: object

Analytic PK moments for independent random equal-rank subspaces.

ambient_dim: int
mean: float
rank: int
variance: float
class transformer_lens.tools.analysis.RankReportRow(idx: int, sigma: float, sigma_ratio: float, is_degenerate: bool, is_null: bool, block_id: int)

Bases: object

One singular direction’s summary, aligned with column idx of U/V.

idx

Position of the direction, matching the column index in U and V.

Type:

int

sigma

The singular value for this direction.

Type:

float

sigma_ratio

sigma normalized by the largest singular value, in [0, 1].

Type:

float

is_degenerate

True when this direction is not attributable on its own: it shares a block with a neighbour (defined only up to a rotation within that block) or it is numerically null.

Type:

bool

is_null

True when sigma_ratio falls below null_rtol, so the singular vector is an arbitrary direction from the map’s null space.

Type:

bool

block_id

Index of the contiguous block this direction belongs to.

Type:

int

block_id: int
idx: int
is_degenerate: bool
is_null: bool
sigma: float
sigma_ratio: float
class transformer_lens.tools.analysis.SparseProbeControl(supports: Int[Tensor, 'repeat selected_feature'], accuracy: Float[Tensor, 'repeat'], precision: Float[Tensor, 'repeat'], recall: Float[Tensor, 'repeat'], f1: Float[Tensor, 'repeat'], roc_auc: Float[Tensor, 'repeat'], average_precision: Float[Tensor, 'repeat'])

Bases: object

Raw held-out metric distributions for one control at one sparsity.

Rows cover only converged repeats, so the row count can be below the request; rejected repeats live in SparseProbeSweep.rejections.

accuracy: Float[Tensor, 'repeat']
average_precision: Float[Tensor, 'repeat']
f1: Float[Tensor, 'repeat']
precision: Float[Tensor, 'repeat']
recall: Float[Tensor, 'repeat']
roc_auc: Float[Tensor, 'repeat']
supports: Int[Tensor, 'repeat selected_feature']
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.SparseProbeMetrics(true_positives: int, true_negatives: int, false_positives: int, false_negatives: int, accuracy: float, precision: float, recall: float, f1: float, roc_auc: float, average_precision: float)

Bases: object

Held-out binary-classification metrics and confusion counts.

Confusion counts, accuracy, precision, recall, and F1 are read at the logit-zero threshold. roc_auc and average_precision are threshold-free and tie-aware: held-out examples with equal logits share one rank and one threshold. The stratified split always holds out both classes; if a class is absent, roc_auc is NaN, and average_precision is NaN without positives and 1.0 without negatives.

accuracy: float
average_precision: float
f1: float
false_negatives: int
false_positives: int
precision: float
recall: float
roc_auc: float
true_negatives: int
true_positives: int
class transformer_lens.tools.analysis.SparseProbeRejection(arm: Literal['random_coordinate', 'label_shuffle'], k: int, repeat: int, support: Int[Tensor, 'selected_feature'], reason: str)

Bases: object

A control fit excluded from the sweep because it did not converge.

Recorded instead of aborting so completed main probes and other controls survive.

arm: Literal['random_coordinate', 'label_shuffle']
k: int
reason: str
repeat: int
support: Int[Tensor, 'selected_feature']
class transformer_lens.tools.analysis.SparseProbeResult(feature_scores: Float[Tensor, 'feature'], selected_features: Int[Tensor, 'selected_feature'], coefficients: Float[Tensor, 'selected_feature'], intercept: Float[Tensor, ''], preprocess_mean: Float[Tensor, 'selected_feature'], preprocess_scale: Float[Tensor, 'selected_feature'], constant_features: Bool[Tensor, 'selected_feature'], train_indices: Int[Tensor, 'train_example'], test_indices: Int[Tensor, 'test_example'], metrics: SparseProbeMetrics, positive_label: int, negative_label: int, train_positive_count: int, train_negative_count: int, test_positive_count: int, test_negative_count: int, preprocess: Literal['none', 'standardize'], std_floor: float, class_weight: Literal['balanced'] | None, l2_strength: float, test_fraction: float, seed: int, k: int, max_iter: int, max_refinement_steps: int, decrement_tolerance: float, objective: float, gradient_inf_norm: float, newton_decrement: float, refinement_steps: int, iterations: int, function_evaluations: int, stop_reason: str)

Bases: object

Result of one train/test k-sparse binary probe fit.

All tensors are detached CPU tensors. Floating-point tensors use float64; coefficients align with selected_features.

class_weight: Literal['balanced'] | None
coefficients: Float[Tensor, 'selected_feature']
constant_features: Bool[Tensor, 'selected_feature']
decrement_tolerance: float
feature_scores: Float[Tensor, 'feature']
function_evaluations: int
gradient_inf_norm: float
intercept: Float[Tensor, '']
iterations: int
k: int
l2_strength: float
max_iter: int
max_refinement_steps: int
metrics: SparseProbeMetrics
negative_label: int
newton_decrement: float
objective: float
positive_label: int
preprocess: Literal['none', 'standardize']
preprocess_mean: Float[Tensor, 'selected_feature']
preprocess_scale: Float[Tensor, 'selected_feature']
refinement_steps: int
seed: int
selected_features: Int[Tensor, 'selected_feature']
std_floor: float
stop_reason: str
test_fraction: float
test_indices: Int[Tensor, 'test_example']
test_negative_count: int
test_positive_count: int
train_indices: Int[Tensor, 'train_example']
train_negative_count: int
train_positive_count: int
class transformer_lens.tools.analysis.SparseProbeSweep(ks: tuple[int, ...], results: tuple[SparseProbeResult, ...], random_coordinate_controls: tuple[SparseProbeControl, ...], label_shuffle_controls: tuple[SparseProbeControl, ...], seed: int, rejections: tuple[SparseProbeRejection, ...] = ())

Bases: object

Probe results and aligned controls over a strictly increasing k-grid.

ks: tuple[int, ...]
label_shuffle_controls: tuple[SparseProbeControl, ...]
random_coordinate_controls: tuple[SparseProbeControl, ...]
rejections: tuple[SparseProbeRejection, ...] = ()
results: tuple[SparseProbeResult, ...]
seed: int
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.VocabularyRanking(values: Float[Tensor, '*leading k'], indices: Int[Tensor, '*leading k'])

Bases: object

Owned CPU copies of signed vocabulary rankings with shape [..., k].

values preserves the floating dtype and sign of logits; indices has dtype torch.int64. Both tensors are detached.

indices: Int[Tensor, '*leading k']
values: Float[Tensor, '*leading k']
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.

transformer_lens.tools.analysis.decompose_head(model, layer: int, head: int, *, which: Sequence[str] = ('QK', 'OV'), eps: float = 0.01, null_rtol: float | None = None) → HeadDecomposition

Decompose a head’s QK (W_Q W_K^T) and/or OV (W_V W_O) maps via SVD.

Weight-space only: this reads the head’s per-block weights via the bridge’s model.blocks[layer].attn accessors and needs no forward pass and no compatibility mode. The returned factors are detached from the model. Whether the model’s weights carry a folded final LayerNorm is recorded on each returned HeadSVD so the readout and patch consumers can refuse a decomposition taken under a different state.

Parameters:

  • model – A TransformerBridge.

  • layer – Layer of the head to decompose.

  • head – Head index within the layer.

  • which – Which maps to decompose, a non-empty sequence drawn from ("QK", "OV"). A bare string is rejected rather than iterated.

  • eps – Relative gap at which a block of adjacent directions ends; directions closer than this share a block.

  • null_rtol – Relative-to-top-singular-value tolerance below which a direction counts as numerically null. Defaults to None, which resolves to d_model * torch.finfo(S.dtype).eps per map, matching torch.linalg.matrix_rank()’s default tolerance.

Returns:

A HeadDecomposition whose QK/OV fields hold a HeadSVD for each requested map.

Raises:

ValueError – If layer or head is out of range, or which is empty, a bare string, or contains an unknown entry.

transformer_lens.tools.analysis.direct_logit_attribution(model, input: str | List[str] | Tensor | None = None, answer_tokens: str | int | Tensor | None = None, incorrect_tokens: str | int | Tensor | None = None, *, unit: str = 'component', pos: int | integer | Tuple[int] | Tuple[int, int] | Tuple[int, int, int] | List[int] | Tensor | ndarray | None = -1, cache: ActivationCache | None = None) → DirectLogitAttribution

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.

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 or with an active learned final-norm weight or bias.

  • NotImplementedError – If a TransformerBridge reports a hybrid block layout (Mamba/SSM/Mixer/LinearAttention).

transformer_lens.tools.analysis.estimate_occupancy(x: Tensor, dictionary: Tensor, *, max_atoms: int = 25, num_control_dictionaries: int = 32, seed: int = 0) → JSpaceOccupancy

Estimate how many dictionary atoms are meaningfully active in x.

Runs the projection-residual recurrence described in _greedy_captured_variance_gains() for exactly max_atoms steps and compares the real per-step captured-variance curve against the same recurrence on num_control_dictionaries random unit-norm dictionaries of the same size. This shares sparse decomposition’s per-step correlation rule, but uses an unconstrained span-projection residual rather than a nonnegative coefficient-fit residual, so their supports need not match. The occupancy is the step of maximum separation between the real and (averaged) control cumulative captured variance – the point past which further atoms add no more than random directions would. Deterministic given seed and needs no threshold. (Captured variance is a projection, hence scale-free, so the random control atoms are simply unit-norm.)

Parameters:

  • x – Target vector, shape [d_model].

  • dictionary – Atom matrix, shape [num_atoms, d_model] (rows are atoms).

  • max_atoms – Number of atoms to select in the real and control recurrences.

  • num_control_dictionaries – Number of random control dictionaries to average over.

  • seed – Seed for the random control dictionaries (reproducibility).

Returns:

An JSpaceOccupancy.

Raises:

ValueError – On complex inputs, a non-2-D dictionary, a target whose length does not match d_model, max_atoms outside [1, num_atoms], num_control_dictionaries < 1, a target with non-finite entries or a non-finite or zero norm, or a dictionary with non-finite or zero-norm atoms.

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, 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.

  • positive_label – Label defining the positive class and score sign.

  • 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 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']

Sweep direct path patches from one source head to all downstream heads.

For every destination head B = (dst_layer, dst_head) where dst_layer > src_layer, patch the contribution of source head A = (src_layer, src_head) into B’s query (or key / value) input, and record the patching metric.

The patch is a linear approximation:

delta_resid = clean_A_result - corrupted_A_result [batch, pos, d_model] delta_B_comp = (delta_resid / ln1_scale) @ W_comp[dst_head]

where W_comp is W_Q, W_K, or W_V according to component.

Parameters:

  • model – Any TransformerLens model (e.g. TransformerBridge).

  • corrupted_tokens – Token IDs for the corrupted input, shape [batch, seq_len].

  • clean_cache – Cached activations from the clean (unpatched) run.

  • corrupted_cache – Cached activations from the corrupted run (needed for ln1 scale).

  • patching_metric – A function mapping the model’s logits tensor to a scalar.

  • src_layer – Layer index of the source attention head.

  • src_head – Head index of the source attention head.

  • component – Which input to patch at the destination head — “q” (default), “k”, or “v”.

  • verbose – Whether to show a tqdm progress bar.

Returns:

results – results[dst_layer, dst_head] is the patching metric when the direct path A → B is patched in. Entries for dst_layer <= src_layer are left as 0.0 (no causal path from A to those layers).

Return type:

Float[Tensor, “n_layers n_heads”]

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']

Full sweep: all (src_layer, src_head) → (dst_layer, dst_head) direct paths.

Returns a 4-D tensor of shape [n_layers, n_heads, n_layers, n_heads]. result[sl, sh, dl, dh] = patching metric when head (sl,sh)’s output is patched directly into head (dl,dh)’s query/key/value input.

Entries where dl <= sl are 0 (no causal path).

This runs O(n_layers * n_heads * n_layers * n_heads) forward passes and is intended for small models or targeted sub-sweeps. For large models prefer calling get_act_patch_direct_path per source head.

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.load_artifact(path: Path) → Dict[str, Any]

Reads and validates a frozen artifact produced by serialize_artifact().

Validates that every required top-level key is present, that schema_version matches SCHEMA_VERSION, that protocol_fingerprint matches a fresh fingerprint_manifest() of the loaded protocol_manifest, and that results_fingerprint matches a fresh results_fingerprint() of the loaded result blocks – catching a hand-edited or corrupted artifact rather than trusting the stored fingerprints blindly.

Raises:

ValueError – Naming the missing or mismatched field.

transformer_lens.tools.analysis.logit_lens(model: Any | ActivationCache, input: str | List[str] | Tensor | None = None, *, cache: ActivationCache | None = None, layers: Sequence[int] | None = None, incl_mid: bool = False, positions: int | integer | Tuple[int] | Tuple[int, int] | Tuple[int, int, int] | List[int] | Tensor | ndarray | None = -1, batch_slice: int | integer | Tuple[int] | Tuple[int, int] | Tuple[int, int, int] | List[int] | Tensor | ndarray | None = None, lens: Callable[[Tensor, int], Tensor] | None = None, apply_ln: bool = True, vocab: Sequence[int] | Tensor | None = None, targets: str | int | Sequence[str | int] | Tensor | None = None, return_type: str = 'logits', top_k: int | None = None, chunk_size: int = 64, enable_grad: bool = False) → LogitLensResult

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).

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.

transformer_lens.tools.analysis.logit_readout(model: Any, vectors: Float[Tensor, '... d_model'], *, apply_ln: bool = True, vocab: Sequence[int] | Tensor | None = None, targets: str | int | Sequence[str | int] | Tensor | None = None, return_type: str = 'logits', top_k: int | None = None, chunk_size: int = 64, enable_grad: bool = False) → LogitReadout

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, or a model whose final norm cannot be applied exactly.

  • NotImplementedError – For encoder-decoder and encoder-only bridges.

transformer_lens.tools.analysis.logit_signature(model, head_svd: HeadSVD, direction: int, tokens: int | Sequence[int] | Tensor) → LogitSignature

Signed logit effect of one OV direction’s rank-1 reconstruction on the given tokens.

Pure weight-space computation: reconstructs the head’s OV output along a single singular direction (S[direction] * V[:, direction], never U - see the module docstring) and projects it through W_U restricted to tokens. Runs no forward pass and builds no cache.

Parameters:

  • model – A TransformerBridge with the final LayerNorm folded into W_U; only its W_U is read.

  • head_svd – An OV HeadSVD from decompose_head().

  • direction – Column index of the singular direction to reconstruct.

  • tokens – Token id(s) to read the logit effect for.

Returns:

A LogitSignature with one value per requested token.

Raises:

  • ValueError – If head_svd.which != "OV", if direction is not in [0, rank), if head_svd was decomposed under a different folded-LayerNorm state than model now has, or if model is a TransformerBridge whose weights do not carry a folded final LayerNorm.

  • DegenerateDirectionError – If direction is not attributable alone (see HeadSVD.require_isolated()).

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.patch_along_directions(model, head_svd: HeadSVD, prompt: str | Tensor, metric: Callable[[Tensor], float], *, keep: Sequence[int] | None = None, ablate: Sequence[int] | None = None, threshold: float | None = None, rng: Generator | None = None, n_baseline: int = 32) → PatchResult

Causally validate a claimed OV subfunction by reconstructing the head’s output onto it.

Requires head_svd.which == "OV": this reconstructs the write/output basis V, and QK has no such vector (see the module docstring). Runs the prompt with use_attn_result enabled once unmodified, once with the head’s hook_result slice reconstructed onto span(head_svd.V[:, retained]), and once per random control subspace. Each control subspace is drawn inside the head’s own OV span span(V) (not from the full residual stream, where a width-w random subspace would keep only w/d_model of a head output that itself occupies only rank/d_model of the stream; an in-span control of width w keeps w/rank of it), so a moved metric is compared against the effect of an arbitrary subspace of this head’s output of the same width. The per-draw control delta magnitudes are averaged over n_baseline draws so one lucky or unlucky draw does not decide the gate, and so controls that mix sign do not cancel into a smaller threshold. Averaging narrows but does not remove the seed dependence of the threshold; see the n_baseline entry below. Restores the model’s prior use_attn_result setting afterward.

The gate’s success condition depends on the mode the caller expressed, because keep=S and ablate=complement(S) resolve to the same retained set and a single “moved more than the control” test would answer only the ablate question. See PatchResult.gated.

Parameters:

  • model – A TransformerBridge.

  • head_svd – An OV HeadSVD from decompose_head().

  • prompt – A single prompt: a string or a [1, pos] token tensor.

  • metric – A function from the model’s logits to a scalar.

  • keep – Direction indices to retain; the rest are zeroed. Exactly one of keep/ablate must be given.

  • ablate – Direction indices to zero; the rest are retained.

  • threshold – Explicit gate threshold. Defaults to None, which uses baseline_delta_metric (already a magnitude) instead.

  • rng – Optional generator for the random control subspaces, for reproducibility. Defaults to a generator seeded with _DEFAULT_BASELINE_SEED so a bare call is reproducible rather than drawing from the global RNG.

  • n_baseline – Number of random in-span control subspaces to average the baseline delta over. Must be at least 1. The per-draw control magnitudes are heavy-tailed, so the averaged threshold still varies with the seed and a direction whose delta sits near it can gate either way; pass an explicit rng for a reproducible verdict and raise n_baseline when the verdict is close to the threshold. Each draw costs one forward pass.

Returns:

A PatchResult describing the patched, baseline, and original metrics.

Raises:

  • ValueError – If head_svd.which != "OV", if head_svd was decomposed under a different folded-LayerNorm state than model now has, if keep/ablate are both given or both omitted, if any index is out of [0, rank), if the retained set is empty (keep=[] or ablate over the full rank) or spans the full rank (keep over every direction) and no explicit threshold is supplied, if n_baseline < 1, or if the model is in training mode.

  • NotImplementedError – If the model’s attention adapter exposes no per-head result, so set_use_attn_result(True) cannot fork the attention output.

  • DegenerateDirectionError – If the retained directions split a degenerate block (see _validate_retained_blocks()).

transformer_lens.tools.analysis.project_activations(model, head_svd: HeadSVD, prompt: str | Tensor) → ActivationProjection

Project a head’s actual per-position output onto its OV singular directions.

Requires head_svd.which == "OV": this projects onto the write/output basis V, and QK has no such vector (see the module docstring). Runs a real forward pass with use_attn_result enabled to read the per-head output (hook_result), then projects it onto head_svd.V. Restores the model’s prior use_attn_result setting afterward, since flipping that config flag as a side effect of a read-only analysis call would surprise a caller who already had hooks or a cache built around its prior state.

Parameters:

  • model – A TransformerBridge.

  • head_svd – An OV HeadSVD from decompose_head().

  • prompt – A single prompt (not a batch): a string or a [pos] or [1, pos] token tensor.

Returns:

An ActivationProjection with the per-position coefficients.

Raises:

  • ValueError – If head_svd.which != "OV", if head_svd was decomposed under a different folded-LayerNorm state than model now has, if prompt is not a single prompt (a tensor that is not 1-D, or is 2-D with a leading dimension greater than one), or if the model is in training mode.

  • NotImplementedError – If the model’s attention adapter exposes no per-head result, so set_use_attn_result(True) cannot fork the attention output.

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_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.run_causal_swap_benchmark(lens: JacobianLens, model: Any, corpus: BenchmarkCorpus, layers: Sequence[int], **trial_kwargs: Any) → Tuple[List[TrialResult], List[BaselineRecord]]

Runs the full causal-swap sweep: baseline filtering, then every surviving trial.

Each (function, source) prompt’s baseline is computed once – it does not depend on layer – and only prompts that survive filter_baseline_capable() proceed to run_causal_swap_trial(), once per remaining layer. Each trial gets its own fresh decomposition_cache: the cache key is (layer, batch_idx, position), which collides across different prompts run as independent single-example forward passes, so a cache may only be reused within one trial’s real/control pair, never across trials.

Parameters:

  • lens – The fitted lens.

  • model – The model to run trials against.

  • corpus – The prompt corpus to sweep.

  • layers – Layers to sweep as an independent trial dimension.

  • **trial_kwargs – Forwarded to run_causal_swap_trial() (control_tolerance, control_seeds, alpha, k; decomposition_cache is not accepted here since each trial always uses its own).

Returns:

(trials, excluded_baselines).

transformer_lens.tools.analysis.select_displacement_matched_control_token(dictionary: Tensor, source_token_id: int, target_token_id: int, excluded_ids: Container[int], active_support: Container[int], *, tolerance: float = 0.1, seed: int = 0) → int | None

Deterministically selects a displacement-matched control token id.

The real condition perturbs the activation by c_src * (a_target - a_source), so the edit’s size is set by ||a_target - a_source||. Matching atom norms leaves that size free: two atoms of equal norm can sit at very different distances from a_source, and the control edit then runs at a different magnitude from the real one. A candidate token id t therefore qualifies when abs(||a_t - a_source|| - ||a_target - a_source||) <= tolerance * ||a_target - a_source|| and t is not source_token_id, not target_token_id, not a member of excluded_ids, and not a member of active_support – an atom already carrying the source coordinate is not a clean control.

One qualifying candidate is picked with a seeded CPU torch.Generator so the same seed always yields the same control token on every device. Returns None when no candidate qualifies; the tolerance is never silently widened and selection never falls back to the globally nearest atom. Returning None rather than raising lets the caller record an empty pool as its own skip status instead of aborting the sweep.

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.

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.solve_coordinate_patch_positions(activations: Tensor, dictionary: Tensor, position_labels: Sequence[int], source_idx: int, target_idx: int, *, layer: 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') → Tuple[Tensor, Dict[Tuple[int, int, int], CoordinatePatch]]

Apply solve_coordinate_patch() independently to every (batch, position) pair.

activations holds one already-sliced forward-pass chunk, shape [batch, num_positions, d_model]; position_labels names the real sequence position of each of its num_positions columns (their order, not their value, need not match column index – the labels only key decomposition_cache and the returned patch dict). Each (batch_idx, position) pair gets its own sparse decomposition and its own CoordinatePatch: a source active in one pair never affects another, matching the anchored, per-pair contract of the underlying primitive.

Parameters:

  • activations – Chunk of activations, shape [batch, num_positions, d_model].

  • dictionary – Atom matrix for this layer, shape [num_atoms, d_model].

  • position_labels – Real sequence position for each column of activations; must have length activations.shape[1].

  • source_idx – Atom index that must occur in the active support of every pair.

  • target_idx – Distinct atom index to receive or exchange the source coordinate.

  • layer – Layer identifier folded into every decomposition_cache key.

  • decomposition_cache – Optional caller-owned mapping from (layer, batch_idx, position) to a previously computed JSpaceDecomposition. A hit skips get_sparse_decomposition() and reuses the strict-compatibility validation already performed inside solve_coordinate_patch(); a miss solves once and stores the result before use.

  • k – Sparse-solver upper bound on a cache miss.

  • mode – "substitute" or "swap", forwarded to solve_coordinate_patch().

  • alpha – Finite interpolation strength, forwarded to solve_coordinate_patch().

  • algorithm – Sparse coefficient-update rule on a cache miss.

Returns:

patched has the same shape as activations, with every (batch_idx, position) entry replaced by that pair’s CoordinatePatch.patched; patches maps (layer, batch_idx, position) to the full CoordinatePatch for that pair.

Return type:

A tuple (patched, patches)

Raises:

ValueError – If activations is not 3-D, if position_labels length does not match activations.shape[1], or if source_idx is not in the active support for any (batch_idx, position) pair – the whole call fails rather than silently skipping that pair.

transformer_lens.tools.analysis.success_rate_ci(successes: Sequence[bool], *, confidence: float = 0.95) → SuccessRateInterval

Computes an exact Clopper-Pearson confidence interval for a success rate.

A percentile bootstrap cannot express uncertainty about an all-failure or all-success sample: every resample of an all-zero vector is all zeros, so the interval collapses to [0, 0] and reports the same certainty for one trial as for a thousand. The exact interval instead widens as the trial count shrinks, giving [0.0, 0.459] for zero successes out of six.

Deterministic by construction: there is no resampling and no seed.

Raises:

ValueError – If successes is empty.

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, 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.

  • positive_label – Label defining the positive class and score sign.

  • 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, 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.vocab_readout(model, head_svd: HeadSVD, *, k: int = 10) → Float[Tensor, 'd_vocab k']

Project the top-k OV output directions through the unembedding.

Requires head_svd.which == "OV": QK produces no write direction to project (see the module docstring). On a TransformerBridge, the final LayerNorm must be folded into W_U, which requires the bridge to have processed its weights with fold_ln enabled on an adapter that supports folding.

Does not call head_svd.require_isolated: a degenerate direction’s vocab readout is still a well-defined projection, unlike a per-direction causal claim, so it is not gated here. The contract that no direction is reported without a passing causal patch is enforced by patch_along_directions().

Parameters:

  • model – A TransformerBridge with the final LayerNorm folded into W_U; only its W_U is read.

  • head_svd – An OV HeadSVD from decompose_head().

  • k – Number of top singular directions to project.

Returns:

column i is direction i’s projection through the unembedding.

Return type:

W_U.T @ head_svd.V[:, :k], shape [d_vocab, k]

Raises:

ValueError – If head_svd.which != "OV", if k is not in (0, rank], if head_svd was decomposed under a different folded-LayerNorm state than model now has, or if model is a TransformerBridge whose weights do not carry a folded final LayerNorm.