From f7863cc94c62dc136e2c4d0c916804d2daeafb00 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 5 Sep 2026 08:51:56 +0530 Subject: [PATCH 01/10] feat(jacobian_lens_coordinate_patch): add per-position patch loop core - Add solve_coordinate_patch_positions: model-free loop applying solve_coordinate_patch independently to every (batch, position) pair in a [batch, num_positions, d_model] chunk - Support an optional caller-owned decomposition_cache keyed (layer, batch_idx, position); a hit skips get_sparse_decomposition, a miss solves once and stores - Fail fast (no try/except) when a source is inactive at any pair, so no partial write reaches the activation tensor - Validate 3-D activations and matching position_labels length - Add unit tests for offline parity, batch independence, fail-fast, cache miss/hit behavior, and label-length validation --- .cursor/rules/transformerlens.mdc | 21 --- docs/source/index.md | 1 + tests/mps/test_mps_basic.py | 19 +++ .../test_jacobian_lens_coordinate_patch.py | 120 ++++++++++++++++++ transformer_lens/tools/analysis/__init__.py | 15 +++ .../jacobian_lens_coordinate_patch.py | 95 +++++++++++++- 6 files changed, 249 insertions(+), 22 deletions(-) delete mode 100644 .cursor/rules/transformerlens.mdc diff --git a/.cursor/rules/transformerlens.mdc b/.cursor/rules/transformerlens.mdc deleted file mode 100644 index 845d076da..000000000 --- a/.cursor/rules/transformerlens.mdc +++ /dev/null @@ -1,21 +0,0 @@ ---- -description: TransformerLens project conventions for Cursor agents. -alwaysApply: true ---- - -Read `AGENTS.md` at the repo root before doing any work. It is the single source of truth for project conventions, quickstart commands, repo layout, hook-naming rules, the HookedTransformer ↔ TransformerBridge mirroring rule, PR conventions, and hard rules. - -Sub-folder `AGENTS.md` files apply when you're working in those directories — read them too: - -- `tests/AGENTS.md` — tier placement, conftest hierarchy, MPS rules -- `transformer_lens/model_bridge/supported_architectures/AGENTS.md` — adapter contract, starter-adapter table, 4-place registration -- `transformer_lens/tools/model_registry/AGENTS.md` — `verify_models` workflow, the `main_benchmark` trap - -Quick reminders that override common defaults: - -- Use `uv`, not `pip` or `poetry`. Install with `uv sync`; run commands with `uv run …` or `make` targets. -- This repo has two parallel systems (`HookedTransformer` legacy and `TransformerBridge` v3). Changes to HookedTransformer that have equivalents in TransformerBridge must be mirrored to TransformerBridge. -- Base PRs against `dev`, not `main`. Never name a branch `main` or `dev`. -- No pre-commit hook is installed. Run `make format` and `uv run mypy .` manually before push. -- Source `.env` (e.g. `set -a; source .env; set +a`) before any HuggingFace-Hub-hitting command. -- Never add `# type: ignore`, never dismiss failing tests as "pre-existing", never add platform skips to dodge CI, never claim drift is "fp noise" without empirical evidence. diff --git a/docs/source/index.md b/docs/source/index.md index 565bb91cf..83aeac699 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -58,6 +58,7 @@ content/hook_system content/compatibility_mode content/ssm_interpretability content/projection_kernel +content/sparse_probing content/jacobian_lens_fitting generated/demos/Jacobian_Lens_Decomposition_Demo content/backward_lens diff --git a/tests/mps/test_mps_basic.py b/tests/mps/test_mps_basic.py index bcf85709c..8de8ccdc7 100644 --- a/tests/mps/test_mps_basic.py +++ b/tests/mps/test_mps_basic.py @@ -28,6 +28,7 @@ SubspaceBasis, projection_kernel, ) +from transformer_lens.tools.analysis.sparse_probing import fit_sparse_probe # Skip the entire module on non-MPS runners (Linux CI, CPU-only Macs) pytestmark = pytest.mark.skipif( @@ -201,6 +202,24 @@ def test_mps_projection_kernel_principal_angles(): _cleanup() +def test_mps_sparse_probe_moves_selected_data_before_float64_conversion(): + """Sparse probing reduces scores on MPS and fits selected columns on CPU.""" + try: + generator = torch.Generator().manual_seed(0) + labels = torch.arange(80) % 2 + features = torch.randn(80, 8, generator=generator) + features[:, 2] += 2 * (2 * labels - 1) + + result = fit_sparse_probe(features.to("mps"), labels, k=2, seed=3) + + assert result.selected_features[0].item() == 2 + assert result.coefficients.device.type == "cpu" + assert result.coefficients.dtype == torch.float64 + assert result.metrics.f1 > 0.9 + finally: + _cleanup() + + def test_mps_softmax_and_layernorm(): """Softmax and LayerNorm — core transformer ops — work on MPS.""" x = torch.randn(4, 16, 64, device="mps", dtype=torch.float32) diff --git a/tests/unit/tools/test_jacobian_lens_coordinate_patch.py b/tests/unit/tools/test_jacobian_lens_coordinate_patch.py index cb3e0c317..61b990ba6 100644 --- a/tests/unit/tools/test_jacobian_lens_coordinate_patch.py +++ b/tests/unit/tools/test_jacobian_lens_coordinate_patch.py @@ -10,6 +10,7 @@ from transformer_lens.tools.analysis.jacobian_lens_coordinate_patch import ( CoordinatePatch, solve_coordinate_patch, + solve_coordinate_patch_positions, ) from transformer_lens.tools.analysis.jacobian_lens_decomposition import ( JSpaceDecomposition, @@ -558,3 +559,122 @@ def test_high_condition_decomposition_is_accepted() -> None: result = solve_coordinate_patch(activation, dictionary, 0, 3, decomposition=decomposition) assert torch.isfinite(result.patched).all() + + +def test_solve_coordinate_patch_positions_matches_offline_solve_coordinate_patch() -> None: + dictionary = torch.eye(3) + activation = torch.tensor([2.0, 5.0, 0.0]) + batched = activation.view(1, 1, 3) + + patched, patches = solve_coordinate_patch_positions( + batched, dictionary, position_labels=[4], source_idx=0, target_idx=1, layer=3, k=2 + ) + expected = solve_coordinate_patch(activation, dictionary, 0, 1, k=2) + + torch.testing.assert_close(patched[0, 0], expected.patched) + torch.testing.assert_close(patches[(3, 0, 4)].coordinates_after, expected.coordinates_after) + + +def test_solve_coordinate_patch_positions_batch_independence() -> None: + dictionary = torch.eye(3) + row0 = torch.tensor([2.0, 5.0, 0.0]) + row1 = torch.tensor([3.0, 1.0, 0.0]) + batched = torch.stack([row0, row1]).unsqueeze(1) + + batched_patched, _ = solve_coordinate_patch_positions( + batched, dictionary, [0], 0, 1, layer=0, k=2 + ) + solo_patched, _ = solve_coordinate_patch_positions( + row0.view(1, 1, 3), dictionary, [0], 0, 1, layer=0, k=2 + ) + + torch.testing.assert_close(batched_patched[0, 0], solo_patched[0, 0]) + + +def test_solve_coordinate_patch_positions_fails_fast_when_any_pair_source_inactive() -> None: + dictionary = torch.eye(3) + activations = torch.stack( + [ + torch.tensor([2.0, 5.0, 0.0]), # source atom 0 active + torch.tensor([0.0, 0.0, 4.0]), # source atom 0 inactive + ] + ).unsqueeze(1) + + with pytest.raises(ValueError, match="active support"): + solve_coordinate_patch_positions(activations, dictionary, [0], 0, 1, layer=0, k=2) + + +def test_solve_coordinate_patch_positions_cache_miss_solves_and_populates( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import transformer_lens.tools.analysis.jacobian_lens_coordinate_patch as core_module + + dictionary = torch.eye(3) + activations = torch.tensor([2.0, 5.0, 0.0]).view(1, 1, 3) + cache: dict = {} + calls = [] + original = core_module.get_sparse_decomposition + + def spy(*args, **kwargs): + calls.append(1) + return original(*args, **kwargs) + + monkeypatch.setattr(core_module, "get_sparse_decomposition", spy) + solve_coordinate_patch_positions( + activations, dictionary, [0], 0, 1, layer=0, decomposition_cache=cache, k=2 + ) + assert len(calls) == 1 + assert (0, 0, 0) in cache + + +def test_solve_coordinate_patch_positions_cache_hit_skips_resolve( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import transformer_lens.tools.analysis.jacobian_lens_coordinate_patch as core_module + + dictionary = torch.eye(3) + activations = torch.tensor([2.0, 5.0, 0.0]).view(1, 1, 3) + cache: dict = {} + calls = [] + original = core_module.get_sparse_decomposition + + def spy(*args, **kwargs): + calls.append(1) + return original(*args, **kwargs) + + monkeypatch.setattr(core_module, "get_sparse_decomposition", spy) + solve_coordinate_patch_positions( + activations, dictionary, [0], 0, 1, layer=0, decomposition_cache=cache, k=2 + ) + assert len(calls) == 1 + solve_coordinate_patch_positions( + activations, dictionary, [0], 0, 1, layer=0, decomposition_cache=cache, k=2 + ) + assert len(calls) == 1 # second call is a pure cache hit + + +def test_solve_coordinate_patch_positions_cache_hit_and_miss_produce_identical_patch() -> None: + dictionary = torch.eye(3) + activations = torch.tensor([2.0, 5.0, 0.0]).view(1, 1, 3) + + fresh_patched, fresh_patches = solve_coordinate_patch_positions( + activations, dictionary, [0], 0, 1, layer=0, k=2 + ) + cache: dict = {} + solve_coordinate_patch_positions( + activations, dictionary, [0], 0, 1, layer=0, decomposition_cache=cache, k=2 + ) + cached_patched, cached_patches = solve_coordinate_patch_positions( + activations, dictionary, [0], 0, 1, layer=0, decomposition_cache=cache, k=2 + ) + + torch.testing.assert_close(fresh_patched, cached_patched) + torch.testing.assert_close(fresh_patches[(0, 0, 0)].patched, cached_patches[(0, 0, 0)].patched) + + +def test_solve_coordinate_patch_positions_rejects_mismatched_position_labels() -> None: + dictionary = torch.eye(3) + activations = torch.tensor([2.0, 5.0, 0.0]).view(1, 1, 3) + + with pytest.raises(ValueError, match="position_labels"): + solve_coordinate_patch_positions(activations, dictionary, [0, 1], 0, 1, layer=0, k=2) diff --git a/transformer_lens/tools/analysis/__init__.py b/transformer_lens/tools/analysis/__init__.py index 496043062..e799fdb43 100644 --- a/transformer_lens/tools/analysis/__init__.py +++ b/transformer_lens/tools/analysis/__init__.py @@ -17,6 +17,7 @@ anchored coordinate patching. - projection_kernel: Basis-invariant subspace overlap and TransformerBridge attention-head OQ/OK/OV affinity. + - sparse_probing: Leakage-safe k-sparse binary probes over activation tensors. """ from transformer_lens.tools.analysis.backward_lens import ( @@ -64,6 +65,14 @@ projection_kernel, random_projection_kernel_moments, ) +from transformer_lens.tools.analysis.sparse_probing import ( + SparseProbeControl, + SparseProbeMetrics, + SparseProbeResult, + SparseProbeSweep, + fit_sparse_probe, + sweep_sparse_probe, +) __all__ = [ "AttentionHeadRef", @@ -84,12 +93,17 @@ "ProjectedFactor", "ProjectionKernelResult", "RandomSubspaceReference", + "SparseProbeControl", + "SparseProbeMetrics", + "SparseProbeResult", + "SparseProbeSweep", "SubspaceBasis", "VocabularyRanking", "WeightLayout", "attention_head_subspace_affinity", "direct_logit_attribution", "estimate_occupancy", + "fit_sparse_probe", "get_act_patch_direct_path", "get_act_patch_direct_path_all_sources", "get_sparse_decomposition", @@ -97,4 +111,5 @@ "projection_kernel", "random_projection_kernel_moments", "solve_coordinate_patch", + "sweep_sparse_probe", ] diff --git a/transformer_lens/tools/analysis/jacobian_lens_coordinate_patch.py b/transformer_lens/tools/analysis/jacobian_lens_coordinate_patch.py index 547905f0d..300bf4dd6 100644 --- a/transformer_lens/tools/analysis/jacobian_lens_coordinate_patch.py +++ b/transformer_lens/tools/analysis/jacobian_lens_coordinate_patch.py @@ -6,7 +6,7 @@ import warnings from dataclasses import dataclass from numbers import Integral, Real -from typing import Optional +from typing import Dict, MutableMapping, Optional, Sequence, Tuple import torch @@ -491,6 +491,99 @@ def solve_coordinate_patch( ) +def solve_coordinate_patch_positions( + activations: torch.Tensor, + dictionary: torch.Tensor, + position_labels: Sequence[int], + source_idx: int, + target_idx: int, + *, + layer: int, + decomposition_cache: Optional[MutableMapping[Tuple[int, int, int], JSpaceDecomposition]] = None, + k: int = DEFAULT_K, + mode: str = "substitute", + alpha: float = 1.0, + algorithm: str = "nonnegative_orthogonal_matching_pursuit", +) -> Tuple[torch.Tensor, Dict[Tuple[int, int, int], CoordinatePatch]]: + """Apply :func:`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 :class:`CoordinatePatch`: a + source active in one pair never affects another, matching the anchored, per-pair contract of + the underlying primitive. + + Args: + 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 :class:`JSpaceDecomposition`. A hit skips + :func:`get_sparse_decomposition` and reuses the strict-compatibility validation + already performed inside :func:`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 :func:`solve_coordinate_patch`. + alpha: Finite interpolation strength, forwarded to :func:`solve_coordinate_patch`. + algorithm: Sparse coefficient-update rule on a cache miss. + + Returns: + A tuple ``(patched, patches)``: ``patched`` has the same shape as ``activations``, with + every ``(batch_idx, position)`` entry replaced by that pair's + :attr:`CoordinatePatch.patched`; ``patches`` maps ``(layer, batch_idx, position)`` to the + full :class:`CoordinatePatch` for that pair. + + 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. + """ + if activations.ndim != 3: + raise ValueError( + "activations must be 3-D [batch, num_positions, d_model], got shape " + f"{tuple(activations.shape)}" + ) + batch_size, num_positions, _ = activations.shape + if len(position_labels) != num_positions: + raise ValueError( + f"position_labels has {len(position_labels)} entries, expected {num_positions} " + "to match activations.shape[1]" + ) + patched = activations.new_empty(activations.shape) + patches: Dict[Tuple[int, int, int], CoordinatePatch] = {} + for batch_idx in range(batch_size): + for column, position in enumerate(position_labels): + key = (layer, batch_idx, position) + x = activations[batch_idx, column] + if decomposition_cache is not None and key in decomposition_cache: + decomposition = decomposition_cache[key] + else: + decomposition = get_sparse_decomposition(x, dictionary, k, algorithm=algorithm) + if decomposition_cache is not None: + decomposition_cache[key] = decomposition + patch = solve_coordinate_patch( + x, + dictionary, + source_idx, + target_idx, + decomposition=decomposition, + k=k, + mode=mode, + alpha=alpha, + algorithm=algorithm, + ) + patched[batch_idx, column] = patch.patched + patches[key] = patch + return patched, patches + + def _postcondition_witnesses( coordinates_before: torch.Tensor, coordinates_after: torch.Tensor, From 764c89589f8b707403dbb4b7922be0885454dfe5 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 5 Sep 2026 17:48:20 +0530 Subject: [PATCH 02/10] test(tools): extract shared J-lens toy-bridge fixtures into conftest - move D_MODEL/N_LAYERS/D_VOCAB/SEQ_LEN/SKIP_FIRST/CORPUS, _ToyBlock, _CausalSumBlock, _ToyTokenizer, _ToyBridge, _NotABridge, _lens, and the toy_model fixture from test_jacobian_lens.py into conftest.py - import the shared symbols back into test_jacobian_lens.py and remove the now-unused contextmanager and HookPoint imports - centralize the shared test setup so additional hook tests can reuse the same fixtures without duplicating roughly 90 lines of test code --- tests/unit/tools/conftest.py | 230 +++++++++++++++++++++++++ tests/unit/tools/test_jacobian_lens.py | 230 ++----------------------- 2 files changed, 242 insertions(+), 218 deletions(-) create mode 100644 tests/unit/tools/conftest.py diff --git a/tests/unit/tools/conftest.py b/tests/unit/tools/conftest.py new file mode 100644 index 000000000..122a91f8b --- /dev/null +++ b/tests/unit/tools/conftest.py @@ -0,0 +1,230 @@ +"""Shared fixtures for the JacobianLens unit-test suite.""" + +from contextlib import contextmanager +from types import SimpleNamespace +from typing import Any + +import pytest +import torch +import torch.nn as nn + +from transformer_lens.hook_points import HookPoint +from transformer_lens.model_bridge import TransformerBridge +from transformer_lens.tools.analysis import JacobianLens +from transformer_lens.utilities.activation_functions import apply_softcap + +D_MODEL = 6 +N_LAYERS = 4 +D_VOCAB = 11 +SEQ_LEN = 9 +SKIP_FIRST = 2 +CORPUS = "unit-test-corpus" + + +class _ToyBlock(nn.Module): + def __init__(self, d_model: int, layer: int, dtype: torch.dtype): + super().__init__() + self.linear = nn.Linear(d_model, d_model, bias=False, dtype=dtype) + nn.init.normal_(self.linear.weight, std=0.2) + self.hook_out = HookPoint() + self.hook_out.name = f"blocks.{layer}.hook_out" + + def forward(self, residual: torch.Tensor) -> torch.Tensor: + return self.hook_out(residual + self.linear(residual)) + + +class _CausalSumBlock(nn.Module): + """Causal cross-position mixing with an exact triangular Jacobian.""" + + def __init__(self, layer: int): + super().__init__() + self.hook_out = HookPoint() + self.hook_out.name = f"blocks.{layer}.hook_out" + + def forward(self, residual: torch.Tensor) -> torch.Tensor: + return self.hook_out(residual.cumsum(dim=1)) + + +class _ToyTokenizer: + def decode(self, token_ids: list[int]) -> str: + return f"token-{token_ids[0]}" + + +class _ToyBridge(TransformerBridge): + """Small real ``TransformerBridge`` subclass with Bridge-native hooks. + + The production constructor needs a Hugging Face model and architecture + adapter. Unit tests only need its public analysis surface, so this subclass + initializes ``nn.Module`` directly while retaining the concrete + ``TransformerBridge`` isinstance contract. + """ + + def __init__( + self, + *, + dtype: torch.dtype = torch.float32, + causal_final_block: bool = False, + ) -> None: + nn.Module.__init__(self) + torch.manual_seed(0) + self.cfg = SimpleNamespace( + n_layers=N_LAYERS, + d_model=D_MODEL, + d_vocab=D_VOCAB, + d_vocab_out=D_VOCAB, + normalization_type="LN", + output_logits_soft_cap=None, + model_name="toy-bridge", + dtype=dtype, + device="cpu", + ) + self.adapter = SimpleNamespace( + supports_generation=True, + get_component_mapping=lambda: { + "blocks": SimpleNamespace(hook_out_is_single_residual_stream=True), + "ln_final": object(), + "unembed": object(), + }, + validate_output_logits_transform=lambda: None, + apply_output_logits_transform=lambda logits: apply_softcap( + logits, self.cfg.output_logits_soft_cap + ), + ) + self.compatibility_mode = False + self._weights_processed = False + self.tokenizer = _ToyTokenizer() + self.embed = nn.Embedding(D_VOCAB, D_MODEL, dtype=dtype) + blocks: list[nn.Module] = [_ToyBlock(D_MODEL, layer, dtype) for layer in range(N_LAYERS)] + if causal_final_block: + blocks[-1] = _CausalSumBlock(N_LAYERS - 1) + self.blocks = nn.ModuleList(blocks) + self.ln_final = nn.Identity() + self.unembed = nn.Linear(D_MODEL, D_VOCAB, bias=False, dtype=dtype) + self.eval() + + @property + def W_U(self) -> torch.Tensor: + return self.unembed.weight.T + + @property + def hook_dict(self) -> dict[str, HookPoint]: + return { + f"blocks.{layer}.hook_out": block.hook_out for layer, block in enumerate(self.blocks) + } + + def parameters(self, recurse: bool = True): + # A production bridge delegates this to its wrapped HF model. This toy + # owns its small modules directly, so enumerate the nn.Module tree. + return nn.Module.parameters(self, recurse=recurse) + + def named_parameters( + self, + prefix: str = "", + recurse: bool = True, + remove_duplicate: bool = True, + ): + return nn.Module.named_parameters( + self, + prefix=prefix, + recurse=recurse, + remove_duplicate=remove_duplicate, + ) + + def to_tokens(self, prompt: str) -> torch.Tensor: + ids = [(3 * index + len(prompt)) % D_VOCAB for index in range(SEQ_LEN)] + return torch.tensor([ids], dtype=torch.long) + + def to_single_token(self, string: str) -> int: + return len(string) % D_VOCAB + + def forward( + self, tokens: torch.Tensor, return_type: str | None = "logits" + ) -> torch.Tensor | None: + residual = self.embed(tokens) + for block in self.blocks: + residual = block(residual) + if return_type is None: + return None + return self.unembed(self.ln_final(residual)) + + @contextmanager + def hooks( + self, + fwd_hooks: list[tuple[str, Any]] = [], + bwd_hooks: list[tuple[str, Any]] = [], + reset_hooks_end: bool = True, + clear_contexts: bool = False, + ): + del clear_contexts + added: list[tuple[HookPoint, str, Any]] = [] + for direction, hook_specs in (("fwd", fwd_hooks), ("bwd", bwd_hooks)): + for name, hook_fn in hook_specs: + hook_point = self.hook_dict[name] + hook_point.add_hook(hook_fn, dir=direction) + handles = hook_point.fwd_hooks if direction == "fwd" else hook_point.bwd_hooks + added.append((hook_point, direction, handles[-1])) + try: + yield self + finally: + if reset_hooks_end: + for hook_point, direction, handle in added: + handle.hook.remove() + handles = hook_point.fwd_hooks if direction == "fwd" else hook_point.bwd_hooks + if handle in handles: + handles.remove(handle) + + def run_with_cache( + self, + input: torch.Tensor, + return_cache_object: bool = False, + remove_batch_dim: bool = False, + names_filter: Any = None, + **kwargs: Any, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + del return_cache_object, remove_batch_dim, kwargs + + def wanted(name: str) -> bool: + if names_filter is None: + return True + if isinstance(names_filter, str): + return name == names_filter + if callable(names_filter): + return bool(names_filter(name)) + return name in names_filter + + cache: dict[str, torch.Tensor] = {} + + def cache_hook(activation: torch.Tensor, hook: HookPoint) -> torch.Tensor: + assert hook.name is not None + cache[hook.name] = activation.detach() + return activation + + cache_hooks = [(name, cache_hook) for name in self.hook_dict if wanted(name)] + with self.hooks(fwd_hooks=cache_hooks): + logits = self(input) + assert logits is not None + return logits, cache + + +class _NotABridge(nn.Module): + def __init__(self) -> None: + super().__init__() + self.cfg = SimpleNamespace(n_layers=N_LAYERS, d_model=D_MODEL) + + +def _lens( + *, + n_prompts: int = 1, + metadata: dict[str, Any] | None = None, +) -> JacobianLens: + return JacobianLens( + {0: torch.eye(D_MODEL)}, + n_prompts=n_prompts, + d_model=D_MODEL, + metadata=metadata, + ) + + +@pytest.fixture(scope="module") +def toy_model() -> _ToyBridge: + return _ToyBridge() diff --git a/tests/unit/tools/test_jacobian_lens.py b/tests/unit/tools/test_jacobian_lens.py index 12229cc2e..877f1d460 100644 --- a/tests/unit/tools/test_jacobian_lens.py +++ b/tests/unit/tools/test_jacobian_lens.py @@ -1,6 +1,5 @@ """Unit tests for the TransformerBridge-only Jacobian lens implementation.""" -from contextlib import contextmanager from enum import IntEnum from inspect import Parameter, signature from types import SimpleNamespace @@ -12,7 +11,18 @@ import torch.nn as nn import transformer_lens.tools.analysis.jacobian_lens as jacobian_lens_module -from transformer_lens.hook_points import HookPoint +from tests.unit.tools.conftest import ( + CORPUS, + D_MODEL, + D_VOCAB, + N_LAYERS, + SEQ_LEN, + SKIP_FIRST, + _lens, + _NotABridge, + _ToyBlock, + _ToyBridge, +) from transformer_lens.model_bridge import TransformerBridge from transformer_lens.model_bridge.generalized_components import AltUpBlockBridge from transformer_lens.model_bridge.supported_architectures.deepseek_v4 import ( @@ -29,13 +39,6 @@ ) from transformer_lens.utilities.activation_functions import apply_softcap -D_MODEL = 6 -N_LAYERS = 4 -D_VOCAB = 11 -SEQ_LEN = 9 -SKIP_FIRST = 2 -CORPUS = "unit-test-corpus" - class _UnsafeMetadataEnum(IntEnum): VALUE = 7 @@ -45,197 +48,6 @@ class _UnsafeMetadataKey(str): pass -class _ToyBlock(nn.Module): - def __init__(self, d_model: int, layer: int, dtype: torch.dtype): - super().__init__() - self.linear = nn.Linear(d_model, d_model, bias=False, dtype=dtype) - nn.init.normal_(self.linear.weight, std=0.2) - self.hook_out = HookPoint() - self.hook_out.name = f"blocks.{layer}.hook_out" - - def forward(self, residual: torch.Tensor) -> torch.Tensor: - return self.hook_out(residual + self.linear(residual)) - - -class _CausalSumBlock(nn.Module): - """Causal cross-position mixing with an exact triangular Jacobian.""" - - def __init__(self, layer: int): - super().__init__() - self.hook_out = HookPoint() - self.hook_out.name = f"blocks.{layer}.hook_out" - - def forward(self, residual: torch.Tensor) -> torch.Tensor: - return self.hook_out(residual.cumsum(dim=1)) - - -class _ToyTokenizer: - def decode(self, token_ids: list[int]) -> str: - return f"token-{token_ids[0]}" - - -class _ToyBridge(TransformerBridge): - """Small real ``TransformerBridge`` subclass with Bridge-native hooks. - - The production constructor needs a Hugging Face model and architecture - adapter. Unit tests only need its public analysis surface, so this subclass - initializes ``nn.Module`` directly while retaining the concrete - ``TransformerBridge`` isinstance contract. - """ - - def __init__( - self, - *, - dtype: torch.dtype = torch.float32, - causal_final_block: bool = False, - ) -> None: - nn.Module.__init__(self) - torch.manual_seed(0) - self.cfg = SimpleNamespace( - n_layers=N_LAYERS, - d_model=D_MODEL, - d_vocab=D_VOCAB, - d_vocab_out=D_VOCAB, - normalization_type="LN", - output_logits_soft_cap=None, - model_name="toy-bridge", - dtype=dtype, - device="cpu", - ) - self.adapter = SimpleNamespace( - supports_generation=True, - get_component_mapping=lambda: { - "blocks": SimpleNamespace(hook_out_is_single_residual_stream=True), - "ln_final": object(), - "unembed": object(), - }, - validate_output_logits_transform=lambda: None, - apply_output_logits_transform=lambda logits: apply_softcap( - logits, self.cfg.output_logits_soft_cap - ), - ) - self.compatibility_mode = False - self._weights_processed = False - self.tokenizer = _ToyTokenizer() - self.embed = nn.Embedding(D_VOCAB, D_MODEL, dtype=dtype) - blocks: list[nn.Module] = [_ToyBlock(D_MODEL, layer, dtype) for layer in range(N_LAYERS)] - if causal_final_block: - blocks[-1] = _CausalSumBlock(N_LAYERS - 1) - self.blocks = nn.ModuleList(blocks) - self.ln_final = nn.Identity() - self.unembed = nn.Linear(D_MODEL, D_VOCAB, bias=False, dtype=dtype) - self.eval() - - @property - def W_U(self) -> torch.Tensor: - return self.unembed.weight.T - - @property - def hook_dict(self) -> dict[str, HookPoint]: - return { - f"blocks.{layer}.hook_out": block.hook_out for layer, block in enumerate(self.blocks) - } - - def parameters(self, recurse: bool = True): - # A production bridge delegates this to its wrapped HF model. This toy - # owns its small modules directly, so enumerate the nn.Module tree. - return nn.Module.parameters(self, recurse=recurse) - - def named_parameters( - self, - prefix: str = "", - recurse: bool = True, - remove_duplicate: bool = True, - ): - return nn.Module.named_parameters( - self, - prefix=prefix, - recurse=recurse, - remove_duplicate=remove_duplicate, - ) - - def to_tokens(self, prompt: str) -> torch.Tensor: - ids = [(3 * index + len(prompt)) % D_VOCAB for index in range(SEQ_LEN)] - return torch.tensor([ids], dtype=torch.long) - - def to_single_token(self, string: str) -> int: - return len(string) % D_VOCAB - - def forward( - self, tokens: torch.Tensor, return_type: str | None = "logits" - ) -> torch.Tensor | None: - residual = self.embed(tokens) - for block in self.blocks: - residual = block(residual) - if return_type is None: - return None - return self.unembed(self.ln_final(residual)) - - @contextmanager - def hooks( - self, - fwd_hooks: list[tuple[str, Any]] = [], - bwd_hooks: list[tuple[str, Any]] = [], - reset_hooks_end: bool = True, - clear_contexts: bool = False, - ): - del clear_contexts - added: list[tuple[HookPoint, str, Any]] = [] - for direction, hook_specs in (("fwd", fwd_hooks), ("bwd", bwd_hooks)): - for name, hook_fn in hook_specs: - hook_point = self.hook_dict[name] - hook_point.add_hook(hook_fn, dir=direction) - handles = hook_point.fwd_hooks if direction == "fwd" else hook_point.bwd_hooks - added.append((hook_point, direction, handles[-1])) - try: - yield self - finally: - if reset_hooks_end: - for hook_point, direction, handle in added: - handle.hook.remove() - handles = hook_point.fwd_hooks if direction == "fwd" else hook_point.bwd_hooks - if handle in handles: - handles.remove(handle) - - def run_with_cache( - self, - input: torch.Tensor, - return_cache_object: bool = False, - remove_batch_dim: bool = False, - names_filter: Any = None, - **kwargs: Any, - ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: - del return_cache_object, remove_batch_dim, kwargs - - def wanted(name: str) -> bool: - if names_filter is None: - return True - if isinstance(names_filter, str): - return name == names_filter - if callable(names_filter): - return bool(names_filter(name)) - return name in names_filter - - cache: dict[str, torch.Tensor] = {} - - def cache_hook(activation: torch.Tensor, hook: HookPoint) -> torch.Tensor: - assert hook.name is not None - cache[hook.name] = activation.detach() - return activation - - cache_hooks = [(name, cache_hook) for name in self.hook_dict if wanted(name)] - with self.hooks(fwd_hooks=cache_hooks): - logits = self(input) - assert logits is not None - return logits, cache - - -class _NotABridge(nn.Module): - def __init__(self) -> None: - super().__init__() - self.cfg = SimpleNamespace(n_layers=N_LAYERS, d_model=D_MODEL) - - def _closed_form_jacobian(model: _ToyBridge, layer: int) -> torch.Tensor: """Exact d h_final / d h_layer for the position-wise linear toy.""" jacobian = torch.eye(D_MODEL) @@ -246,24 +58,6 @@ def _closed_form_jacobian(model: _ToyBridge, layer: int) -> torch.Tensor: return jacobian -def _lens( - *, - n_prompts: int = 1, - metadata: dict[str, Any] | None = None, -) -> JacobianLens: - return JacobianLens( - {0: torch.eye(D_MODEL)}, - n_prompts=n_prompts, - d_model=D_MODEL, - metadata=metadata, - ) - - -@pytest.fixture(scope="module") -def toy_model() -> _ToyBridge: - return _ToyBridge() - - @pytest.fixture(scope="module") def fitted_lens(toy_model: _ToyBridge) -> JacobianLens: return JacobianLens.fit( From 83f1d477421ecae1a7db3dfecc95357edd72991f Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 5 Sep 2026 18:09:12 +0530 Subject: [PATCH 03/10] feat(jacobian_lens): expose coordinate_patch_hooks on the Bridge - Add JacobianLens.coordinate_patch_hooks, a forward-hook variant of the offline coordinate_patch primitive, following the swap_hooks builder pattern - Solve one J-space coordinate patch per (batch, position) pair at each layer via solve_coordinate_patch_positions, with an optional caller-owned decomposition_cache keyed (layer, batch_idx, position) - Require positions explicitly and reject identical source/target tokens; fail fast on any inactive source rather than partially patching a batch - Warn once per call naming the layer x position count that performs a live vocabulary-scale solve on every cache miss - Add a dedicated test file covering shape parity, warning-once, cache hit/miss, oracle parity with offline coordinate_patch, and uncaught error and warning propagation through the hook --- ...st_jacobian_lens_coordinate_patch_hooks.py | 186 ++++++++++++++++++ .../tools/analysis/jacobian_lens.py | 126 +++++++++++- 2 files changed, 311 insertions(+), 1 deletion(-) create mode 100644 tests/unit/tools/test_jacobian_lens_coordinate_patch_hooks.py diff --git a/tests/unit/tools/test_jacobian_lens_coordinate_patch_hooks.py b/tests/unit/tools/test_jacobian_lens_coordinate_patch_hooks.py new file mode 100644 index 000000000..207fd18aa --- /dev/null +++ b/tests/unit/tools/test_jacobian_lens_coordinate_patch_hooks.py @@ -0,0 +1,186 @@ +"""Hook-based tests for dynamic J-space coordinate patching.""" + +import warnings +from typing import Any + +import pytest +import torch + +from tests.unit.tools.conftest import D_MODEL, D_VOCAB, _lens, _ToyBridge +from transformer_lens.tools.analysis import JacobianLens + +# The toy vocab dictionary has only D_VOCAB atoms, fewer than the library DEFAULT_K, so every +# real solve in this file passes an explicit k within [1, D_VOCAB]. +SOLVE_K = 8 + + +def _active_source_and_distinct_target( + lens: JacobianLens, model: _ToyBridge, prompt: str, layer: int, position: int +) -> tuple[int, int]: + """Discover a real active source at (layer, position) instead of guessing a token id -- + coordinate_patch_hooks requires the source to be active, and a toy model's real activations + are not hand-computable in advance.""" + decomposition = lens.decompose(model, prompt, layer=layer, position=position, k=SOLVE_K) + source_id = int(decomposition.support[0]) + target_id = (source_id + 1) % D_VOCAB + return source_id, target_id + + +def test_coordinate_patch_hooks_shape_mirrors_swap_hooks(toy_model: _ToyBridge) -> None: + with pytest.warns(UserWarning): + hooks = _lens().coordinate_patch_hooks(toy_model, 3, 5, layers=[0], positions=[0]) + assert [name for name, _ in hooks] == ["blocks.0.hook_out"] + assert callable(hooks[0][1]) + + +def test_coordinate_patch_hooks_warns_once_naming_layer_and_position_counts( + toy_model: _ToyBridge, +) -> None: + lens = JacobianLens( + {0: torch.eye(D_MODEL), 1: torch.eye(D_MODEL)}, n_prompts=1, d_model=D_MODEL + ) + with pytest.warns(UserWarning, match=r"2 layer\(s\) x 2 position\(s\)") as record: + lens.coordinate_patch_hooks(toy_model, 3, 5, layers=[0, 1], positions=[0, 1]) + assert len(record) == 1 + + +def test_coordinate_patch_hooks_rejects_unfitted_layer(toy_model: _ToyBridge) -> None: + with pytest.raises(ValueError, match="source layers"): + _lens().coordinate_patch_hooks(toy_model, 3, 5, layers=[2], positions=[0]) + + +def test_coordinate_patch_hooks_requires_nonempty_positions(toy_model: _ToyBridge) -> None: + with pytest.raises(ValueError, match="positions"): + _lens().coordinate_patch_hooks(toy_model, 3, 5, layers=[0], positions=[]) + + +def test_coordinate_patch_hooks_rejects_identical_tokens(toy_model: _ToyBridge) -> None: + with pytest.raises(ValueError, match="same token|identical|distinct"): + _lens().coordinate_patch_hooks(toy_model, 3, 3, layers=[0], positions=[0]) + + +def test_coordinate_patch_hooks_changes_only_requested_positions(toy_model: _ToyBridge) -> None: + lens = _lens() + prompt = "a toy prompt" + tokens = toy_model.to_tokens(prompt) + _, baseline = toy_model.run_with_cache(tokens) + source_id, target_id = _active_source_and_distinct_target(lens, toy_model, prompt, 0, -1) + + with pytest.warns(UserWarning): + hooks = lens.coordinate_patch_hooks( + toy_model, source_id, target_id, layers=[0], positions=[-1], k=SOLVE_K + ) + with toy_model.hooks(fwd_hooks=hooks): + _, patched = toy_model.run_with_cache(tokens) + + delta = patched["blocks.0.hook_out"] - baseline["blocks.0.hook_out"] + torch.testing.assert_close(delta[:, :-1], torch.zeros_like(delta[:, :-1])) + + +def test_coordinate_patch_hooks_oracle_parity_with_offline_coordinate_patch( + toy_model: _ToyBridge, +) -> None: + lens = _lens() + prompt = "a toy prompt" + tokens = toy_model.to_tokens(prompt) + _, baseline = toy_model.run_with_cache(tokens) + pre_hook_activation = baseline["blocks.0.hook_out"][0, -1, :].float() + source_id, target_id = _active_source_and_distinct_target(lens, toy_model, prompt, 0, -1) + + with pytest.warns(UserWarning): + hooks = lens.coordinate_patch_hooks( + toy_model, source_id, target_id, layers=[0], positions=[-1], k=SOLVE_K + ) + with toy_model.hooks(fwd_hooks=hooks): + _, patched = toy_model.run_with_cache(tokens) + + expected = lens.coordinate_patch( + toy_model, + pre_hook_activation, + layer=0, + source_token=source_id, + target_token=target_id, + k=SOLVE_K, + ) + torch.testing.assert_close(patched["blocks.0.hook_out"][0, -1, :].float(), expected.patched) + + +def test_decomposition_cache_hit_skips_resolve_across_hook_firings( + toy_model: _ToyBridge, monkeypatch: pytest.MonkeyPatch +) -> None: + import transformer_lens.tools.analysis.jacobian_lens_coordinate_patch as core_module + + lens = _lens() + prompt = "a toy prompt" + tokens = toy_model.to_tokens(prompt) + source_id, target_id = _active_source_and_distinct_target(lens, toy_model, prompt, 0, -1) + cache: dict = {} + calls = [] + original = core_module.get_sparse_decomposition + + def spy(*args: Any, **kwargs: Any) -> Any: + calls.append(1) + return original(*args, **kwargs) + + monkeypatch.setattr(core_module, "get_sparse_decomposition", spy) + with pytest.warns(UserWarning): + hooks = lens.coordinate_patch_hooks( + toy_model, + source_id, + target_id, + layers=[0], + positions=[-1], + decomposition_cache=cache, + k=SOLVE_K, + ) + with toy_model.hooks(fwd_hooks=hooks): + toy_model(tokens) + first_call_count = len(calls) + assert first_call_count >= 1 + + with toy_model.hooks(fwd_hooks=hooks): + toy_model(tokens) + assert len(calls) == first_call_count # second forward pass is entirely cache hits + + +def test_coordinate_patch_hooks_propagates_core_errors_uncaught( + toy_model: _ToyBridge, monkeypatch: pytest.MonkeyPatch +) -> None: + """The fail-fast design decision: an inactive-source error from the core loop must abort the + whole forward pass, not be caught and turned into a partial/silent patch.""" + import transformer_lens.tools.analysis.jacobian_lens as jacobian_lens_module + + def raise_inactive_source(*args: Any, **kwargs: Any) -> Any: + raise ValueError("source_idx=3 is not in the decomposition's active support") + + monkeypatch.setattr( + jacobian_lens_module, "solve_coordinate_patch_positions", raise_inactive_source + ) + with pytest.warns(UserWarning): + hooks = _lens().coordinate_patch_hooks(toy_model, 3, 5, layers=[0], positions=[0]) + with pytest.raises(ValueError, match="active support"): + with toy_model.hooks(fwd_hooks=hooks): + toy_model(toy_model.to_tokens("a toy prompt")) + + +def test_coordinate_patch_hooks_warnings_propagate_uncaught( + toy_model: _ToyBridge, monkeypatch: pytest.MonkeyPatch +) -> None: + """Warnings raised by the core loop (conditioning, near-parallel) must reach the caller through + the hook -- not be swallowed or re-wrapped.""" + import transformer_lens.tools.analysis.jacobian_lens as jacobian_lens_module + + def fake_solve( + activations: torch.Tensor, dictionary, position_labels, source_idx, target_idx, **kwargs + ): + warnings.warn( + "coordinate-patch source and target atoms are near-parallel (stub)", UserWarning + ) + return activations.clone(), {} + + monkeypatch.setattr(jacobian_lens_module, "solve_coordinate_patch_positions", fake_solve) + with pytest.warns(UserWarning, match="coordinate_patch_hooks"): + hooks = _lens().coordinate_patch_hooks(toy_model, 3, 5, layers=[0], positions=[0]) + with pytest.warns(UserWarning, match="near-parallel"): + with toy_model.hooks(fwd_hooks=hooks): + toy_model(toy_model.to_tokens("a toy prompt")) diff --git a/transformer_lens/tools/analysis/jacobian_lens.py b/transformer_lens/tools/analysis/jacobian_lens.py index 400cc0841..4c061402d 100644 --- a/transformer_lens/tools/analysis/jacobian_lens.py +++ b/transformer_lens/tools/analysis/jacobian_lens.py @@ -63,7 +63,17 @@ import warnings from dataclasses import dataclass from importlib.metadata import version -from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union +from typing import ( + Any, + Callable, + Dict, + List, + MutableMapping, + Optional, + Sequence, + Tuple, + Union, +) import torch from jaxtyping import Float, Int @@ -72,6 +82,7 @@ from transformer_lens.tools.analysis.jacobian_lens_coordinate_patch import ( CoordinatePatch, solve_coordinate_patch, + solve_coordinate_patch_positions, ) from transformer_lens.tools.analysis.jacobian_lens_decomposition import ( DEFAULT_K, @@ -1375,6 +1386,119 @@ def transform( ) return hooks + def coordinate_patch_hooks( + self, + model: Any, + source_token: TokenInput, + target_token: TokenInput, + layers: Sequence[int], + *, + positions: Sequence[int], + decomposition_cache: Optional[ + MutableMapping[Tuple[int, int, int], JSpaceDecomposition] + ] = None, + k: int = DEFAULT_K, + 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 :meth:`coordinate_patch`, which edits one already-captured activation offline, + this installs a forward hook that solves :func:`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. + + Args: + 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 + :class:`~transformer_lens.tools.analysis.jacobian_lens_decomposition.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 inactive at any ``(batch_idx, + position)`` pair touched by a hook firing (the whole forward pass fails rather + than silently patching a subset). + + Warns: + UserWarning: Once per call, naming the number of layers and positions that will + perform a live vocabulary-scale solve on every cache miss. + """ + self.validate_model(model) + if not positions: + raise ValueError("positions must contain at least one index") + resolved_layers = [_normalize_layer(layer, model.cfg.n_layers) for layer in layers] + source_id, target_id = _to_token_ids(model, [source_token, target_token]) + if source_id == target_id: + raise ValueError( + "source_token and target_token resolve to the same token id; " + "a coordinate patch would be a silent no-op" + ) + warnings.warn( + f"coordinate_patch_hooks installs {len(resolved_layers)} layer(s) x " + f"{len(positions)} position(s) of coordinate-patch hooks; every (batch, position) " + "pair not already present in decomposition_cache performs a vocabulary-scale sparse " + "decomposition on every forward pass", + UserWarning, + stacklevel=2, + ) + + requested = tuple(positions) + hooks = [] + for layer in resolved_layers: + dictionary = self.lens_vector_dictionary(model, layer) + + def hook_fn( + activation: Float[torch.Tensor, "batch pos d_model"], + hook: Any, + layer: int = layer, + dictionary: torch.Tensor = dictionary, + ) -> Float[torch.Tensor, "batch pos d_model"]: + hook_name = getattr(hook, "name", "intervention hook") + _validate_residual_activation( + activation, d_model=model.cfg.d_model, hook_name=hook_name + ) + normalized = _normalize_positions(requested, activation.shape[1]) + selected = activation[:, normalized, :].float().to(dictionary.device) + patched, _ = solve_coordinate_patch_positions( + selected, + dictionary, + normalized, + source_id, + target_id, + layer=layer, + decomposition_cache=decomposition_cache, + k=k, + mode=mode, + alpha=alpha, + algorithm=algorithm, + ) + output = activation.clone() + output[:, normalized, :] = patched.to( + device=activation.device, dtype=activation.dtype + ) + return output + + hooks.append((_resid_post_hook_name(layer), hook_fn)) + return hooks + # ------------------------------------------------------------------ # # fitting # # ------------------------------------------------------------------ # From ae0a0102d36807130fc3edbca18a75eb70a0f651 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 5 Sep 2026 19:33:24 +0530 Subject: [PATCH 04/10] feat(jacobian_lens): export and document coordinate_patch_hooks Export solve_coordinate_patch_positions from tools.analysis and document the distinction between offline and dynamic/hooked coordinate patching. Replace the stale dynamic-patching claim in jacobian_lens_fitting.md and add documentation covering required positions, caller-owned decomposition_cache, per-pair fail-fast behavior, and the once-per-call cost warning. Add a cached GPT-2 integration test verifying alpha=0 is an exact no-op and untouched positions remain bit-identical. --- docs/source/content/jacobian_lens_fitting.md | 50 ++++++++++++++++++-- tests/integration/test_jacobian_lens.py | 44 +++++++++++++++++ transformer_lens/tools/analysis/__init__.py | 4 +- 3 files changed, 94 insertions(+), 4 deletions(-) diff --git a/docs/source/content/jacobian_lens_fitting.md b/docs/source/content/jacobian_lens_fitting.md index 7492226a1..f432f4e4c 100644 --- a/docs/source/content/jacobian_lens_fitting.md +++ b/docs/source/content/jacobian_lens_fitting.md @@ -350,9 +350,53 @@ Read the measured condition number and cosine straight out of each warning messa judge how much to trust the edit. Without `decomposition=`, coordinate patching runs the vocabulary-scale sparse decomposition first. -Reuse a compatible result for repeated edits to avoid that scan. The method returns an offline -activation and diagnostics, not forward hooks; applying patches dynamically would otherwise perform -a vocabulary-scale solve inside the model run. +Reuse a compatible result for repeated edits to avoid that scan. `coordinate_patch` returns an +offline activation and diagnostics, not forward hooks; see "Dynamic coordinate-patch hooks" below +for the live variant, which performs the same scan per `(batch, position)` pair on every forward +pass unless a cache hit avoids it. + +### Dynamic coordinate-patch hooks + +`JacobianLens.coordinate_patch_hooks` installs the same anchored edit as a forward hook, so it can +run inside `model.run_with_hooks(...)` or `model.generate(...)` instead of on one pre-captured +activation: + +```python +hooks = lens.coordinate_patch_hooks( + model, + source_token=source_id, + target_token=" Paris", + layers=[6], + positions=[-1], + mode="substitute", +) +with model.hooks(fwd_hooks=hooks): + patched_logits = model(tokens) +``` + +Two departures from `coordinate_patch`, both deliberate: + +- **`positions` is required.** There is no full-sequence default: each hooked position performs + its own vocabulary-scale sparse decomposition (unless a cache hit avoids it), and a silent + full-sequence default would trigger that scan at every position without the caller asking for + it. +- **`decomposition_cache` (optional, caller-owned).** A plain `dict` (or any `MutableMapping`) + keyed `(layer, batch_idx, position)`. Pass the same dict across repeated `model.hooks(...)` calls + on the same prompt (e.g. holding the prompt fixed while varying `alpha` or `mode` in an + interactive loop) to skip the vocabulary-scale scan on every hit; a miss solves once and + populates the cache. This is purely a performance path — a cache hit and a fresh solve produce + an identical patch. + +Every `(batch_idx, position)` pair gets its own independent decomposition and edit: a source +concept inactive at one pair never affects another pair in the same batch or call. If the source +is inactive at **any** pair touched by a hook firing, the whole forward pass raises — there is no +silent partial application across a batch. + +Calling `coordinate_patch_hooks(...)` emits one `UserWarning` naming the number of layers and +positions it installs, since every one of those `(layer, position)` combinations performs a live, +vocabulary-scale solve on every forward pass unless `decomposition_cache` already has an entry for +it. The conditioning and near-parallel warnings described above still fire from inside the hook, +per pair, exactly as they would from an offline `coordinate_patch` call on that pair's activation. ### Interpreting the numbers honestly diff --git a/tests/integration/test_jacobian_lens.py b/tests/integration/test_jacobian_lens.py index d1e7bb9c8..55197a735 100644 --- a/tests/integration/test_jacobian_lens.py +++ b/tests/integration/test_jacobian_lens.py @@ -518,6 +518,50 @@ def test_coordinate_patch_gpt2_preserves_anchored_frame(published_gpt2_lens, gpt assert torch.equal(no_op.patched, activation) +def test_coordinate_patch_hooks_gpt2_no_op_and_leaves_other_positions_unchanged( + published_gpt2_lens, gpt2_bridge +): + """coordinate_patch_hooks through a real run_with_hooks pass: alpha=0 is an exact no-op and + untouched positions are bit-identical -- algebraic invariants only, no token-flip claim (see + build-plan.md §9's policy against behavior-dependent assertions).""" + layer = 6 + tokens = gpt2_bridge.to_tokens(PROMPT) + hook = f"blocks.{layer}.hook_out" + _, baseline_cache = gpt2_bridge.run_with_cache(tokens, names_filter=lambda name: name == hook) + baseline_activation = baseline_cache[hook][0, -1, :].float() + decomposition = published_gpt2_lens.decompose( + gpt2_bridge, baseline_activation, layer=layer, k=8 + ) + source_id = int(decomposition.support[0]) + dictionary = published_gpt2_lens.lens_vector_dictionary(gpt2_bridge, layer) + units = dictionary / dictionary.norm(dim=1, keepdim=True) + pair_cosines = (units @ units[source_id]).abs() + pair_cosines[source_id] = torch.inf + target_id = int(pair_cosines.argmin().item()) + + with pytest.warns(UserWarning, match="coordinate_patch_hooks"): + no_op_hooks = published_gpt2_lens.coordinate_patch_hooks( + gpt2_bridge, source_id, target_id, layers=[layer], positions=[-1], alpha=0.0 + ) + with gpt2_bridge.hooks(fwd_hooks=no_op_hooks): + _, no_op_cache = gpt2_bridge.run_with_cache(tokens, names_filter=lambda name: name == hook) + torch.testing.assert_close( + no_op_cache[hook][0, -1, :].float(), baseline_activation, atol=1e-5, rtol=1e-5 + ) + + with pytest.warns(UserWarning, match="coordinate_patch_hooks"): + edit_hooks = published_gpt2_lens.coordinate_patch_hooks( + gpt2_bridge, source_id, target_id, layers=[layer], positions=[-1], alpha=0.5 + ) + with gpt2_bridge.hooks(fwd_hooks=edit_hooks): + _, edited_cache = gpt2_bridge.run_with_cache(tokens, names_filter=lambda name: name == hook) + edited = edited_cache[hook].float() + baseline_full = baseline_cache[hook].float() + torch.testing.assert_close(edited[:, :-1, :], baseline_full[:, :-1, :], atol=1e-5, rtol=1e-5) + assert torch.isfinite(edited).all() + assert not torch.allclose(edited[0, -1, :], baseline_full[0, -1, :]) + + def test_occupancy_gpt2_activation_is_a_small_positive_integer(published_gpt2_lens, gpt2_bridge): """occupancy on a real GPT-2 activation returns a positive integer within ``[1, max_atoms]``, with per-step real and control captured-variance curves of the right shape. We assert shape and diff --git a/transformer_lens/tools/analysis/__init__.py b/transformer_lens/tools/analysis/__init__.py index e799fdb43..5500d9b1e 100644 --- a/transformer_lens/tools/analysis/__init__.py +++ b/transformer_lens/tools/analysis/__init__.py @@ -14,7 +14,7 @@ - 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. + anchored coordinate patching (offline and dynamic/hooked). - projection_kernel: Basis-invariant subspace overlap and TransformerBridge attention-head OQ/OK/OV affinity. - sparse_probing: Leakage-safe k-sparse binary probes over activation tensors. @@ -45,6 +45,7 @@ from transformer_lens.tools.analysis.jacobian_lens_coordinate_patch import ( CoordinatePatch, solve_coordinate_patch, + solve_coordinate_patch_positions, ) from transformer_lens.tools.analysis.jacobian_lens_decomposition import ( JSpaceDecomposition, @@ -111,5 +112,6 @@ "projection_kernel", "random_projection_kernel_moments", "solve_coordinate_patch", + "solve_coordinate_patch_positions", "sweep_sparse_probe", ] From 26fc437a7d7bc74bda05fbf62f5c135fe137aa73 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 5 Sep 2026 23:33:57 +0530 Subject: [PATCH 05/10] fix(jacobian_lens): drop leaked sparse_probing exports from analysis __init__ Commit 02ffc277 accidentally added sparse_probing imports and __all__ entries to transformer_lens/tools/analysis/__init__.py without committing the sparse_probing module itself. On CI (which only checks out tracked files) importing transformer_lens.tools.analysis raised ModuleNotFoundError, breaking package import and failing every job that imports transformer_lens (unit, docstring, compatibility, benchmark, coverage, notebooks). Remove the out-of-scope sparse_probing exports; that work belongs to its own PR. --- transformer_lens/tools/analysis/__init__.py | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/transformer_lens/tools/analysis/__init__.py b/transformer_lens/tools/analysis/__init__.py index 5500d9b1e..36db37881 100644 --- a/transformer_lens/tools/analysis/__init__.py +++ b/transformer_lens/tools/analysis/__init__.py @@ -17,7 +17,6 @@ anchored coordinate patching (offline and dynamic/hooked). - projection_kernel: Basis-invariant subspace overlap and TransformerBridge attention-head OQ/OK/OV affinity. - - sparse_probing: Leakage-safe k-sparse binary probes over activation tensors. """ from transformer_lens.tools.analysis.backward_lens import ( @@ -66,14 +65,6 @@ projection_kernel, random_projection_kernel_moments, ) -from transformer_lens.tools.analysis.sparse_probing import ( - SparseProbeControl, - SparseProbeMetrics, - SparseProbeResult, - SparseProbeSweep, - fit_sparse_probe, - sweep_sparse_probe, -) __all__ = [ "AttentionHeadRef", @@ -94,17 +85,12 @@ "ProjectedFactor", "ProjectionKernelResult", "RandomSubspaceReference", - "SparseProbeControl", - "SparseProbeMetrics", - "SparseProbeResult", - "SparseProbeSweep", "SubspaceBasis", "VocabularyRanking", "WeightLayout", "attention_head_subspace_affinity", "direct_logit_attribution", "estimate_occupancy", - "fit_sparse_probe", "get_act_patch_direct_path", "get_act_patch_direct_path_all_sources", "get_sparse_decomposition", @@ -113,5 +99,4 @@ "random_projection_kernel_moments", "solve_coordinate_patch", "solve_coordinate_patch_positions", - "sweep_sparse_probe", ] From 69ece957d9c88f48b1c81ae085956dbe31ec84d6 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 9 Sep 2026 11:17:43 +0530 Subject: [PATCH 06/10] chore(jacobian_lens): drop unrelated sparse-probing and cursor changes from PR scope - Restore .cursor/rules/transformerlens.mdc deleted on this branch (unrelated to coordinate-patch hooks) - Remove sparse_probing import + test from tests/mps/test_mps_basic.py that broke tests/mps collection with ModuleNotFoundError (accidentally picked up from another branch) - Remove content/sparse_probing toctree line from docs/source/index.md referencing a nonexistent page --- .cursor/rules/transformerlens.mdc | 21 +++++++++++++++++++++ docs/source/index.md | 1 - tests/mps/test_mps_basic.py | 19 ------------------- 3 files changed, 21 insertions(+), 20 deletions(-) create mode 100644 .cursor/rules/transformerlens.mdc diff --git a/.cursor/rules/transformerlens.mdc b/.cursor/rules/transformerlens.mdc new file mode 100644 index 000000000..845d076da --- /dev/null +++ b/.cursor/rules/transformerlens.mdc @@ -0,0 +1,21 @@ +--- +description: TransformerLens project conventions for Cursor agents. +alwaysApply: true +--- + +Read `AGENTS.md` at the repo root before doing any work. It is the single source of truth for project conventions, quickstart commands, repo layout, hook-naming rules, the HookedTransformer ↔ TransformerBridge mirroring rule, PR conventions, and hard rules. + +Sub-folder `AGENTS.md` files apply when you're working in those directories — read them too: + +- `tests/AGENTS.md` — tier placement, conftest hierarchy, MPS rules +- `transformer_lens/model_bridge/supported_architectures/AGENTS.md` — adapter contract, starter-adapter table, 4-place registration +- `transformer_lens/tools/model_registry/AGENTS.md` — `verify_models` workflow, the `main_benchmark` trap + +Quick reminders that override common defaults: + +- Use `uv`, not `pip` or `poetry`. Install with `uv sync`; run commands with `uv run …` or `make` targets. +- This repo has two parallel systems (`HookedTransformer` legacy and `TransformerBridge` v3). Changes to HookedTransformer that have equivalents in TransformerBridge must be mirrored to TransformerBridge. +- Base PRs against `dev`, not `main`. Never name a branch `main` or `dev`. +- No pre-commit hook is installed. Run `make format` and `uv run mypy .` manually before push. +- Source `.env` (e.g. `set -a; source .env; set +a`) before any HuggingFace-Hub-hitting command. +- Never add `# type: ignore`, never dismiss failing tests as "pre-existing", never add platform skips to dodge CI, never claim drift is "fp noise" without empirical evidence. diff --git a/docs/source/index.md b/docs/source/index.md index 83aeac699..565bb91cf 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -58,7 +58,6 @@ content/hook_system content/compatibility_mode content/ssm_interpretability content/projection_kernel -content/sparse_probing content/jacobian_lens_fitting generated/demos/Jacobian_Lens_Decomposition_Demo content/backward_lens diff --git a/tests/mps/test_mps_basic.py b/tests/mps/test_mps_basic.py index 8de8ccdc7..bcf85709c 100644 --- a/tests/mps/test_mps_basic.py +++ b/tests/mps/test_mps_basic.py @@ -28,7 +28,6 @@ SubspaceBasis, projection_kernel, ) -from transformer_lens.tools.analysis.sparse_probing import fit_sparse_probe # Skip the entire module on non-MPS runners (Linux CI, CPU-only Macs) pytestmark = pytest.mark.skipif( @@ -202,24 +201,6 @@ def test_mps_projection_kernel_principal_angles(): _cleanup() -def test_mps_sparse_probe_moves_selected_data_before_float64_conversion(): - """Sparse probing reduces scores on MPS and fits selected columns on CPU.""" - try: - generator = torch.Generator().manual_seed(0) - labels = torch.arange(80) % 2 - features = torch.randn(80, 8, generator=generator) - features[:, 2] += 2 * (2 * labels - 1) - - result = fit_sparse_probe(features.to("mps"), labels, k=2, seed=3) - - assert result.selected_features[0].item() == 2 - assert result.coefficients.device.type == "cpu" - assert result.coefficients.dtype == torch.float64 - assert result.metrics.f1 > 0.9 - finally: - _cleanup() - - def test_mps_softmax_and_layernorm(): """Softmax and LayerNorm — core transformer ops — work on MPS.""" x = torch.randn(4, 16, 64, device="mps", dtype=torch.float32) From 90a4c4c4174305fe8365c3333a50e9b182f45b3d Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 9 Sep 2026 11:32:02 +0530 Subject: [PATCH 07/10] fix(jacobian_lens): straight-through gradient in coordinate-patch position loop - Scatter `x + patch.delta` instead of the detached `patch.patched` in solve_coordinate_patch_positions, keeping x's graph; bitwise identical forward for float32 activations. - Add test_solve_coordinate_patch_positions_preserves_gradient_to_patched_position: asserts finite nonzero grad at the patched position (exactly 0 before) and bitwise-equal forward output. --- .../test_jacobian_lens_coordinate_patch.py | 22 +++++++++++++++++++ .../jacobian_lens_coordinate_patch.py | 6 ++++- 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/tests/unit/tools/test_jacobian_lens_coordinate_patch.py b/tests/unit/tools/test_jacobian_lens_coordinate_patch.py index 61b990ba6..081faf3e3 100644 --- a/tests/unit/tools/test_jacobian_lens_coordinate_patch.py +++ b/tests/unit/tools/test_jacobian_lens_coordinate_patch.py @@ -672,6 +672,28 @@ def test_solve_coordinate_patch_positions_cache_hit_and_miss_produce_identical_p torch.testing.assert_close(fresh_patches[(0, 0, 0)].patched, cached_patches[(0, 0, 0)].patched) +def test_solve_coordinate_patch_positions_preserves_gradient_to_patched_position() -> None: + dictionary = torch.eye(3) + activation = torch.tensor([2.0, 5.0, 0.0]) + activations = activation.view(1, 1, 3).clone().requires_grad_(True) + + patched, _ = solve_coordinate_patch_positions( + activations, dictionary, position_labels=[4], source_idx=0, target_idx=1, layer=3, k=2 + ) + + # Forward output is unchanged: the straight-through edit is bitwise identical to ``patch.patched``. + expected = solve_coordinate_patch(activation, dictionary, 0, 1, k=2) + assert torch.equal(patched[0, 0].detach(), expected.patched) + + patched.sum().backward() + assert activations.grad is not None + grad_at_patch = activations.grad[0, 0] + # Straight-through gives an identity gradient at the patched position; the severed + # ``patch.patched`` path left this exactly zero. + assert torch.isfinite(grad_at_patch).all() + assert float(grad_at_patch.abs().sum()) > 0.0 + + def test_solve_coordinate_patch_positions_rejects_mismatched_position_labels() -> None: dictionary = torch.eye(3) activations = torch.tensor([2.0, 5.0, 0.0]).view(1, 1, 3) diff --git a/transformer_lens/tools/analysis/jacobian_lens_coordinate_patch.py b/transformer_lens/tools/analysis/jacobian_lens_coordinate_patch.py index 300bf4dd6..095377ffd 100644 --- a/transformer_lens/tools/analysis/jacobian_lens_coordinate_patch.py +++ b/transformer_lens/tools/analysis/jacobian_lens_coordinate_patch.py @@ -579,7 +579,11 @@ def solve_coordinate_patch_positions( alpha=alpha, algorithm=algorithm, ) - patched[batch_idx, column] = patch.patched + # Straight-through edit: ``patch.patched`` is a detached ``x.float() + delta`` and would + # sever the gradient path at every patched position. ``patch.delta`` is a detached + # constant w.r.t. ``x``, so ``x + patch.delta`` is bitwise identical for the float32 + # activations while keeping ``x``'s graph -- an identity gradient at the patch. + patched[batch_idx, column] = x + patch.delta patches[key] = patch return patched, patches From c408c39e3756e7d3319f70a353f56cac552f40f6 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 9 Sep 2026 11:37:32 +0530 Subject: [PATCH 08/10] test(jacobian_lens): pin batch component of decomposition_cache key - Add test_solve_coordinate_patch_positions_populates_batch_component_of_cache_key: two distinct batch rows share one cache at the same (layer, position); assert both (layer, 0, position) and (layer, 1, position) populate with different active supports. - Closes the gap where the batch-agnostic [1, 1, d_model] fixtures let a collapsed (layer, 0, position) key pass the whole unit tier; the new test fails when the key drops batch_idx. --- .../test_jacobian_lens_coordinate_patch.py | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/tests/unit/tools/test_jacobian_lens_coordinate_patch.py b/tests/unit/tools/test_jacobian_lens_coordinate_patch.py index 081faf3e3..4b847727b 100644 --- a/tests/unit/tools/test_jacobian_lens_coordinate_patch.py +++ b/tests/unit/tools/test_jacobian_lens_coordinate_patch.py @@ -672,6 +672,41 @@ def test_solve_coordinate_patch_positions_cache_hit_and_miss_produce_identical_p torch.testing.assert_close(fresh_patches[(0, 0, 0)].patched, cached_patches[(0, 0, 0)].patched) +def test_solve_coordinate_patch_positions_populates_batch_component_of_cache_key() -> None: + """The batch component of the ``(layer, batch_idx, position)`` cache key is load-bearing. + + Every other cache test uses a ``[1, 1, d_model]`` activation, so collapsing the key to + ``(layer, 0, position)`` would pass the whole tier. Here two distinct batch rows share one + caller-owned cache at the same ``(layer, position)``: each must occupy its own slot and get its + own solve, so a collapsed key -- which would collide row 1 onto row 0's entry -- is caught. + """ + dictionary = torch.eye(3) + row0 = torch.tensor([2.0, 5.0, 0.0]) # active support {0, 1} + row1 = torch.tensor([3.0, 0.0, 4.0]) # active support {0, 2} -- differs from row0 + activations = torch.stack([row0, row1]).unsqueeze(1) # [2, 1, 3] + cache: dict = {} + + solve_coordinate_patch_positions( + activations, + dictionary, + position_labels=[7], + source_idx=0, + target_idx=1, + layer=0, + decomposition_cache=cache, + k=2, + ) + + # Both batch rows are keyed separately at the shared (layer, position); neither collided. + assert (0, 0, 7) in cache + assert (0, 1, 7) in cache + # Each row solved independently: the two decompositions have different active supports, so + # row 1 did not reuse row 0's cache entry. + row0_support = cache[(0, 0, 7)].support.sort().values + row1_support = cache[(0, 1, 7)].support.sort().values + assert not torch.equal(row0_support, row1_support) + + def test_solve_coordinate_patch_positions_preserves_gradient_to_patched_position() -> None: dictionary = torch.eye(3) activation = torch.tensor([2.0, 5.0, 0.0]) From 3655f514a7a3ec9af75b1fd818f37860c46ae4d3 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 9 Sep 2026 11:45:53 +0530 Subject: [PATCH 09/10] docs+test(jacobian_lens): document top-k support precondition and cover plural-layer install - Strengthen coordinate_patch_hooks Raises note: source must be in every patched position's top-k active support after earlier band hooks edit the residual, not merely active on a clean pass; substitute/swap remove the source coordinate, so stacking layers/positions makes it progressively harder. - Add the first plural (two-layer) end-to-end hook test, pinning the per-closure layer=layer/dictionary=dictionary binding via distinct per-layer dictionaries and per-layer cache keys. --- ...st_jacobian_lens_coordinate_patch_hooks.py | 105 +++++++++++++++++- .../tools/analysis/jacobian_lens.py | 12 +- 2 files changed, 113 insertions(+), 4 deletions(-) diff --git a/tests/unit/tools/test_jacobian_lens_coordinate_patch_hooks.py b/tests/unit/tools/test_jacobian_lens_coordinate_patch_hooks.py index 207fd18aa..d393f0f9e 100644 --- a/tests/unit/tools/test_jacobian_lens_coordinate_patch_hooks.py +++ b/tests/unit/tools/test_jacobian_lens_coordinate_patch_hooks.py @@ -6,7 +6,7 @@ import pytest import torch -from tests.unit.tools.conftest import D_MODEL, D_VOCAB, _lens, _ToyBridge +from tests.unit.tools.conftest import D_MODEL, D_VOCAB, SEQ_LEN, _lens, _ToyBridge from transformer_lens.tools.analysis import JacobianLens # The toy vocab dictionary has only D_VOCAB atoms, fewer than the library DEFAULT_K, so every @@ -184,3 +184,106 @@ def fake_solve( with pytest.warns(UserWarning, match="near-parallel"): with toy_model.hooks(fwd_hooks=hooks): toy_model(toy_model.to_tokens("a toy prompt")) + + +def test_coordinate_patch_hooks_plural_install_binds_each_layer_to_its_own_dictionary( + toy_model: _ToyBridge, +) -> None: + """First end-to-end test of a *plural* install: two fitted layers, one real forward pass. + + It pins two things no single-layer test above reaches: + + * The per-closure default-argument binding ``layer=layer, dictionary=dictionary`` in + ``coordinate_patch_hooks``. Without it, Python's late-binding closures make every hook + capture the *last* loop iteration's ``layer``/``dictionary`` -- so layer 0's hook would + solve against layer 1's dictionary and key its cache under layer 1. The two layers are + given deliberately different (non-scalar) dictionaries so a mis-bound hook produces a + numerically wrong edit; coordinate patching is scale-covariant, so a uniform rescale would + leave the edit unchanged and hide the mis-binding. + * The band precondition documented in the ``Raises`` note: the source must stay in the active + support at *both* layers after the earlier hook has already edited the residual. The + fixture searches for a source that satisfies it rather than assuming one does. + """ + lens = JacobianLens( + {0: torch.eye(D_MODEL), 1: torch.diag(torch.linspace(0.5, 2.0, D_MODEL))}, + n_prompts=1, + d_model=D_MODEL, + ) + prompt = "a toy prompt" + tokens = toy_model.to_tokens(prompt) + position = -1 + normalized_position = SEQ_LEN - 1 # -1 over the toy model's fixed sequence length + + _, baseline = toy_model.run_with_cache(tokens) + clean_layer0 = baseline["blocks.0.hook_out"][0, position].float() + + # Find one source that is active at layer 0 *and* still active at layer 1 after layer 0's edit + # -- exactly the band precondition. Candidates come from layer 0's clean active support. + layer0_support = [ + int(token) + for token in lens.decompose( + toy_model, tokens, layer=0, position=position, k=SOLVE_K + ).support + ] + source_id = -1 + target_id = -1 + layer1_input: torch.Tensor | None = None + for candidate in layer0_support: + candidate_target = (candidate + 1) % D_VOCAB + if candidate_target == candidate: + continue + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + layer0_only = lens.coordinate_patch_hooks( + toy_model, candidate, candidate_target, layers=[0], positions=[position], k=SOLVE_K + ) + with toy_model.hooks(fwd_hooks=layer0_only): + _, edited = toy_model.run_with_cache(tokens) + after_edit = edited["blocks.1.hook_out"][0, position].float() + support = [ + int(token) + for token in lens.decompose(toy_model, after_edit, layer=1, k=SOLVE_K).support + ] + if candidate in support: + source_id, target_id, layer1_input = candidate, candidate_target, after_edit + break + assert ( + layer1_input is not None + ), "no source stays active at both layers; fixture needs a new prompt" + + # Install on BOTH fitted layers and run one real forward pass through model.hooks(...). + cache: dict = {} + with pytest.warns(UserWarning, match=r"2 layer\(s\) x 1 position\(s\)"): + hooks = lens.coordinate_patch_hooks( + toy_model, + source_id, + target_id, + layers=[0, 1], + positions=[position], + decomposition_cache=cache, + k=SOLVE_K, + ) + with toy_model.hooks(fwd_hooks=hooks): + _, patched = toy_model.run_with_cache(tokens) + + # Layer 0's hook must solve against layer 0's OWN dictionary. A mis-bound closure would use + # layer 1's dictionary here and produce a different edit of the clean layer-0 activation. + expected_layer0 = lens.coordinate_patch( + toy_model, clean_layer0, layer=0, source_token=source_id, target_token=target_id, k=SOLVE_K + ) + torch.testing.assert_close( + patched["blocks.0.hook_out"][0, position].float(), expected_layer0.patched + ) + + # Layer 1's hook must solve against layer 1's own dictionary, on the residual as edited by + # layer 0 upstream (``layer1_input`` was captured with only layer 0's hook installed). + expected_layer1 = lens.coordinate_patch( + toy_model, layer1_input, layer=1, source_token=source_id, target_token=target_id, k=SOLVE_K + ) + torch.testing.assert_close( + patched["blocks.1.hook_out"][0, position].float(), expected_layer1.patched + ) + + # Each layer keyed its own cache slot; a late-bound ``layer`` would collapse both onto layer 1. + assert (0, 0, normalized_position) in cache + assert (1, 0, normalized_position) in cache diff --git a/transformer_lens/tools/analysis/jacobian_lens.py b/transformer_lens/tools/analysis/jacobian_lens.py index 4c061402d..5f9caf8fa 100644 --- a/transformer_lens/tools/analysis/jacobian_lens.py +++ b/transformer_lens/tools/analysis/jacobian_lens.py @@ -1433,9 +1433,15 @@ def coordinate_patch_hooks( Raises: ValueError: If ``positions`` is empty, if ``source_token`` and ``target_token`` - resolve to the same id, or if ``source_token`` is inactive at any ``(batch_idx, - position)`` pair touched by a hook firing (the whole forward pass fails rather - than silently patching a subset). + 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 From 454b5873174ff5bc85568edc6067da66ba75908a Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 9 Sep 2026 11:49:11 +0530 Subject: [PATCH 10/10] docs(jacobian_lens): correct dynamic-hook generate() claim and cache-reuse caveat - Drop the model.generate(...) claim from the dynamic coordinate-patch hooks section: on GPT-2 with the published lens, generate re-decomposes each freshly generated token whose top-k support generally lacks the source concept, so fail-fast aborts generation. Keep run_with_hooks(...). - Document that the (layer, batch_idx, position) cache key uses the chunk-local position index, so decomposition_cache is valid only across identical-chunking passes; warn against reuse across decode steps (use_past_kv_cache prefill vs per-step shapes collide) or differently-shaped prompts. --- docs/source/content/jacobian_lens_fitting.md | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/docs/source/content/jacobian_lens_fitting.md b/docs/source/content/jacobian_lens_fitting.md index f432f4e4c..9bec2f472 100644 --- a/docs/source/content/jacobian_lens_fitting.md +++ b/docs/source/content/jacobian_lens_fitting.md @@ -358,8 +358,7 @@ pass unless a cache hit avoids it. ### Dynamic coordinate-patch hooks `JacobianLens.coordinate_patch_hooks` installs the same anchored edit as a forward hook, so it can -run inside `model.run_with_hooks(...)` or `model.generate(...)` instead of on one pre-captured -activation: +run inside `model.run_with_hooks(...)` instead of on one pre-captured activation: ```python hooks = lens.coordinate_patch_hooks( @@ -385,7 +384,14 @@ Two departures from `coordinate_patch`, both deliberate: on the same prompt (e.g. holding the prompt fixed while varying `alpha` or `mode` in an interactive loop) to skip the vocabulary-scale scan on every hit; a miss solves once and populates the cache. This is purely a performance path — a cache hit and a fresh solve produce - an identical patch. + an identical patch. The `position` in the `(layer, batch_idx, position)` key is the + **chunk-local** index into the activation the hook sees, not an absolute sequence position, so the + cache is valid **only across passes with identical chunking** — the same prompt sliced the same + way. Do not reuse one cache across decode steps (with `use_past_kv_cache=True` the prefill sees + `[1, seq, d]` while each decode step sees `[1, 1, d]`, so step 2's position `0` collides with the + prefill's position `0` and the mismatched precomputed coordinates raise an NNLS stationarity error + that does not point back here) or across prompts of different lengths. For those, use a fresh cache + per shape. Every `(batch_idx, position)` pair gets its own independent decomposition and edit: a source concept inactive at one pair never affects another pair in the same batch or call. If the source