feat(jacobian_lens): add dynamic J-space coordinate-patch hooks

docs/source/content/jacobian_lens_fitting.md

@@ -350,9 +350,59 @@ 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(...)` 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. 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
+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
 

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

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