transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark module¶

Causal coordinate-swap benchmark for JacobianLens.coordinate_patch_hooks.

Measures whether an anchored J-space coordinate edit installed live inside a forward pass via coordinate_patch_hooks causes a directional change in model output, under three controls: baseline-capability filtering (only intervene on prompts the model already answers correctly), displacement-matched random-atom controls (isolate “this concept mattered” from “any edit of similar magnitude would have mattered”), and exact Clopper-Pearson uncertainty on every reported rate. Each trial draws its control arm under several seeds and records every draw, so the control arm’s own spread is visible rather than folded into a single number.

This module is layered bottom-up: the model-free prompt corpus and rank/margin metric, baseline-capability filtering, displacement-matched control-token selection, a per-trial runner that wires the first three together with real coordinate_patch_hooks calls against a live model, and exact Clopper-Pearson confidence intervals plus a versioned JSON artifact schema that fingerprints both the protocol and the result blocks and embeds the corpus definition. Generating the frozen artifact consumed by demos/Jacobian_Lens_Coordinate_Patch_Benchmark_Demo.ipynb is a separate entry point: transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark_cli.

class transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark.AnswerMetrics(top1_token_id: int, target_rank: int, target_is_top1: bool, target_tied_for_top: bool, target_logit_margin: float)¶

Bases: object

Rank/margin/tie metrics for one target token against one next-token logit vector.

target_is_top1: bool¶

target_logit_margin: float¶

target_rank: int¶

target_tied_for_top: bool¶

top1_token_id: int¶

class transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark.BaselineRecord(function: str, source: str, prompt: str, metrics: AnswerMetrics)¶

Bases: object

A source prompt’s own-answer metrics under the unperturbed baseline forward pass.

function: str¶

metrics: AnswerMetrics¶

prompt: str¶

source: str¶

class transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark.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.jacobian_lens_causal_swap_benchmark.ControlDraw(seed: int, token_id: int, metrics: AnswerMetrics)¶

Bases: object

One seeded control-arm draw: the token it selected and how that condition scored.

metrics: AnswerMetrics¶

seed: int¶

token_id: int¶

class transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark.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.jacobian_lens_causal_swap_benchmark.PromptTrialSpec(function: str, source: str, target: str, prompt: str, source_answer: str, target_answer: str)¶

Bases: object

One (function, ordered source/target concept pair) prompt instance.

function: str¶

prompt: str¶

source: str¶

source_answer: str¶

target: str¶

target_answer: str¶

class transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark.SuccessRateInterval(point_estimate: float, ci_low: float, ci_high: float, n_trials: int, n_successes: int, confidence: float)¶

Bases: object

An exact Clopper-Pearson confidence interval around a success rate.

n_trials and n_successes are the counts the interval was computed from, so a reader can see the denominator behind the bounds rather than inferring it.

ci_high: float¶

ci_low: float¶

confidence: float¶

n_successes: int¶

n_trials: int¶

point_estimate: float¶

class transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark.TrialResult(function: str, source: str, target: str, layer: int, status: Literal['ok', 'skipped_source_inactive', 'skipped_no_control_token'], baseline: AnswerMetrics, real_target_metrics: AnswerMetrics | None, control_token_id: int | None, control_target_metrics: AnswerMetrics | None, control_draws: List[ControlDraw], error: str | None)¶

Bases: object

One (function, source, target, layer) causal-swap trial’s full record.

control_token_id and control_target_metrics mirror the first entry of control_draws so a single-draw reader sees the same fields it always did. The full list carries the control arm’s own spread across seeds, which one draw cannot measure.

baseline: AnswerMetrics¶

control_draws: List[ControlDraw]¶

control_target_metrics: AnswerMetrics | None¶

control_token_id: int | None¶

error: str | None¶

function: str¶

layer: int¶

real_target_metrics: AnswerMetrics | None¶

source: str¶

status: Literal['ok', 'skipped_source_inactive', 'skipped_no_control_token']¶

target: str¶

transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark.build_protocol_manifest(**fields: Any) → Dict[str, Any]¶

Assembles a protocol manifest, requiring the benchmark’s fixed field set.

Requires at least _REQUIRED_MANIFEST_FIELDS, the same key set (and, where they overlap, the same string values) as Jacobian_Lens_Demo.ipynb’s existing protocol_manifest cell. Beyond that cell’s fields the manifest carries the full corpus definition (corpus) and the repo/path/revision triple it was taken from (corpus_repo, corpus_path, corpus_revision), so the prompts and answers that decide capability and scoring enter the fingerprint instead of a bare name.

Raises:: ValueError – If any required field is missing, naming the missing field(s).

transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark.compute_answer_metrics(logits: Tensor, target_token_id: int) → AnswerMetrics¶

Computes rank/margin/tie metrics for target_token_id against logits.

Ports Jacobian_Lens_Demo.ipynb’s _target_metrics cell verbatim (arithmetic unchanged, only renamed and restructured into a dataclass), so results stay directly comparable with that notebook’s already-reviewed success/rank definitions.

transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark.corpus_definition(corpus: BenchmarkCorpus) → Dict[str, Any]¶

Serializes a corpus into the JSON shape embedded in the protocol manifest.

The prompts and answers decide which prompts are baseline-capable and how every condition is scored, so they belong inside the fingerprinted manifest rather than travelling as a bare name. Key order is normalized so the same corpus always serializes identically.

transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark.filter_baseline_capable(baselines: Sequence[BaselineRecord]) → Tuple[List[BaselineRecord], List[BaselineRecord]]¶

Splits baseline records into (capable, excluded) prompts.

A prompt is baseline-capable when the model’s own deterministic argmax already matches the source’s answer (metrics.target_is_top1) and that answer is not tied for the maximum (metrics.target_tied_for_top); only such prompts are eligible for later intervention trials, so an edit’s effect is never measured against a prompt the unperturbed model already gets wrong. A tied-for-top baseline is treated as not capable: argmax returns the lowest index on an exact tie, so admitting ties would let token-id order decide whether a prompt enters the trial set. Order-preserving in both outputs; never mutates baselines.

transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark.fingerprint_manifest(manifest: Dict[str, Any]) → str¶: Fingerprints a protocol manifest with the recipe Jacobian_Lens_Demo.ipynb uses.

transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark.iter_prompt_trials(corpus: BenchmarkCorpus) → Iterator[PromptTrialSpec]¶

Yields one spec per (function, ordered source/target concept pair).

Ordered pairs are every element of itertools.permutations(corpus.concepts, 2), the same cross product Jacobian_Lens_Demo.ipynb’s country benchmark already uses, now as tested library code instead of a notebook cell.

transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark.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.jacobian_lens_causal_swap_benchmark.results_fingerprint(artifact: Mapping[str, Any]) → str¶

Fingerprints an artifact’s result blocks with the manifest recipe.

protocol_fingerprint covers only the manifest, so an artifact whose trials, excluded baselines, or confidence-interval blocks were rewritten would still load clean and print the same provenance hash. Hashing those blocks separately catches that tampering while leaving the manifest recipe and its stored value untouched.

transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark.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.jacobian_lens_causal_swap_benchmark.run_causal_swap_trial(lens: JacobianLens, model: Any, trial_spec: PromptTrialSpec, layer: int, *, decomposition_cache: MutableMapping[Tuple[int, int, int], JSpaceDecomposition] | None = None, control_tolerance: float = 0.1, control_seeds: Sequence[int] = (0,), alpha: float = 1.0, k: int = 25) → TrialResult¶

Runs one causal-swap trial: baseline, then real and control coordinate-patch conditions.

A baseline forward pass scores the prompt’s own (unperturbed) source answer. A displacement-matched control token is then selected from layer’s lens-vector dictionary, excluding the source token, the real target token, both prompts’ answer tokens, and every atom in the layer’s active support. The real and control conditions each install coordinate_patch_hooks at layer and position -1, sharing one decomposition_cache so only the first of the two performs the vocabulary-scale decomposition; both are scored against the target’s own answer token, so a real-vs-control gap isolates “swapping toward this concept mattered” from “any edit of this magnitude would have mattered.”

The control arm is drawn once per entry in control_seeds. A single draw leaves the control arm’s own variance unmeasured, so a reader cannot tell a real-vs-control gap from draw variation; the per-seed results are carried on control_draws. Every draw reuses the same decomposition_cache, so only the first pays the vocabulary-scale solve.

Two conditions produce a skip rather than a result. When the source is not in the layer’s active support, the trial is recorded with status="skipped_source_inactive". When no displacement-matched control token survives the exclusions, the trial is recorded with status="skipped_no_control_token". Both carry a diagnostic message in error.

The active-support precondition is checked up front, against the same decomposition the hooks consume, rather than inferred from a caught exception. coordinate_patch_hooks raises plain ValueError for several unrelated protocol faults (a same-id source and target, an out-of-vocabulary answer, non-finite logits) and defines no narrower subclass to discriminate on, so catching ValueError would file those as skips and drop them from both denominators. Every other error therefore propagates unchanged. coordinate_patch_hooks’s own ``UserWarning``s (both the per-call install notice and any solver-side conditioning warning) are not suppressed here and propagate to the caller unchanged.

Parameters:

lens – The fitted lens.
model – The model to run trials against.
trial_spec – The prompt, source/target concepts, and their answer words.
layer – The single layer to patch at.
decomposition_cache – Shared cache passed to both the real and control coordinate_patch_hooks calls. A fresh cache is used if omitted.
control_tolerance – Relative tolerance for the displacement-matched control token.
control_seeds – Seeds to draw the control arm under. Each seed selects its own control token and contributes one ControlDraw; the first seed’s draw is mirrored onto control_token_id and control_target_metrics.
alpha – Interpolation strength forwarded to coordinate_patch_hooks.
k – Sparse-solver upper bound forwarded to coordinate_patch_hooks.

Returns:

The trial’s TrialResult.

Raises:

ValueError – If the source and target concepts resolve to the same token id, which would make the coordinate patch a silent no-op, or if control_seeds is empty.

transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark.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.jacobian_lens_causal_swap_benchmark.serialize_artifact(manifest: Dict[str, Any], trials: Sequence[TrialResult], excluded_baselines: Sequence[BaselineRecord], real_ci: SuccessRateInterval, control_ci: SuccessRateInterval) → Dict[str, Any]¶

Assembles a JSON-serializable artifact dict from a benchmark run’s results.

n_independent_prompts records how many distinct (function, source) prompts the pooled rate rests on. Pooling trials as independent draws overstates the evidence when several trials share one prompt, so the count travels with the rate.

layers_executed lists the sorted distinct layers that produced an ok trial, and trial_status_counts counts trials per status, so a reader can see how many of the swept layers contributed nothing and how many trials were skipped rather than executed.

transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark.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.