transformer_lens.tools.analysis.jacobian_lens_coordinate_patch module

Anchored coordinate edits over sparse Jacobian-lens decompositions.

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