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

[janmenjayap]

---

Summary

Adds JacobianLens.coordinate_patch_hooks(...), a forward-hook variant of the offline
coordinate_patch(...) primitive shipped in #1741, so an anchored J-space coordinate edit can run
live inside model.run_with_hooks(...) / model.generate(...) instead of only on one
already-captured activation.

Underneath, a new model-free core function, solve_coordinate_patch_positions(...), loops
solve_coordinate_patch(...) over every (batch_idx, position) pair in a [batch, num_positions, d_model] chunk, with an optional caller-owned decomposition cache keyed (layer, batch_idx, position). coordinate_patch_hooks(...) is a thin wrapper that builds one forward hook per layer —
following the same _resid_post_hook_name / eager dictionary-build pattern as swap_hooks — whose
hook body slices the requested positions, calls the core loop, and writes the result back.

This is PR2 of the plan tracked by the companion issue. Arbitrary multi-slot permutations, a
causal-swap benchmark, and HookedTransformer support remain explicitly deferred (PR3/PR4).

Fixes #1748

Motivation

coordinate_patch(...) edits one pre-captured activation offline and returns a report-only
CoordinatePatch. Studying how a patched coordinate propagates through the rest of a forward pass —
or patching during generate(...) — previously required manually re-capturing, patching, and
re-injecting activations outside the hook system. #1739 deferred the hooked variant because each
hooked position would otherwise fire its own vocabulary-scale sparse decomposition live during the
forward pass. This PR ships that variant while making the per-position cost explicit (a required
positions argument and a construction-time warning) rather than hiding it behind a convenient API.

What ships (commit by commit)

1. feat(jacobian_lens_coordinate_patch): add per-position patch loop core
   New model-free solve_coordinate_patch_positions(...) in jacobian_lens_coordinate_patch.py:
   loops solve_coordinate_patch(...) over every (batch_idx, position) pair — never once per
   position shared across the batch, since two batch rows can have different active supports.
   Supports an optional caller-owned decomposition_cache keyed (layer, batch_idx, position); a
   cache hit skips get_sparse_decomposition(...), a miss solves once and populates it. No
   try/except around the per-pair loop — an inactive source at any one pair propagates
   solve_coordinate_patch's existing ValueError out of the whole call, so nothing is a partial
   write. Model-free tests first, mirroring PR1's own core-before-wrapper split.

2. test(tools): extract shared J-lens toy-bridge fixtures into conftest
   Mechanical move, no behavior change. Extracts the _ToyBridge fake TransformerBridge (and its
   supporting constants/classes/fixture) that test_jacobian_lens.py already defined into
   tests/unit/tools/conftest.py, so the new hook-test file in commit 3 can reuse it instead of
   duplicating ~90 lines that would otherwise risk drifting between two copies. test_jacobian_lens.py
   imports from the conftest; its pass/fail set is byte-identical to before the move.

3. feat(jacobian_lens): expose coordinate_patch_hooks on the Bridge
JacobianLens.coordinate_patch_hooks(model, source_token, target_token, layers, *, positions, decomposition_cache=None, ...) resolves token inputs and layers exactly like coordinate_patch,
   then builds one forward hook per layer via solve_coordinate_patch_positions. positions is
   required — no full-sequence default, since a silent default would trigger a vocabulary-scale
   sparse solve at every position. Emits one UserWarning per call, at hook-construction time,
   naming len(layers) x len(positions) — not a per-position warning, and not an attempt to report
   the exact number of live (non-cache-hit) solves, which is only known at forward time. New
   dedicated test file (test_jacobian_lens_coordinate_patch_hooks.py), not folded into the
   already-1600+-line test_jacobian_lens.py.

4. feat(jacobian_lens): export and document coordinate_patch_hooks
   Exports solve_coordinate_patch_positions from transformer_lens.tools.analysis, extends the
   package docstring, replaces the now-stale "dynamic patching would require a vocabulary-scale
   solve" claim in jacobian_lens_fitting.md with a new #### Dynamic coordinate-patch hooks
   subsection, and adds one cached-model GPT-2 integration test asserting algebraic invariants
   only (alpha=0 is an exact no-op; untouched positions are bit-identical through a real
   run_with_hooks pass) — not a token-flip claim.

Design decisions (per the tracking issue)

#	Decision	Status
1	Batch handling: loop (batch_idx, position) pairs explicitly, never once per position shared across the batch	Required correctness, not a style choice — settled by construction
2	Fail-fast (not skip-with-warning) when a source is inactive at any one pair in a batch	Needs maintainer sign-off — the alternative (try/except ValueError + warnings.warn + copy the row through unedited) is a two-line hunk if preferred
3	decomposition_cache is caller-owned (Optional[MutableMapping[Tuple[int, int, int], JSpaceDecomposition]]), not a JacobianLens attribute	Needs maintainer sign-off, lower risk than #2
4	Warning fires once per coordinate_patch_hooks(...) call, at construction time, naming layers x positions counts (not the exact number of live solves)	This plan's own call, not yet reviewed — flagged explicitly here

Please weigh in on #2 and #3 in particular — they're the most likely to get pushback and are cheap
to change if so.

Testing

• tests/unit/tools/test_jacobian_lens_coordinate_patch.py — model-free core: parity with
  solve_coordinate_patch, batch independence, fail-fast on any inactive source, cache miss/hit
  behavior, cache-hit vs. fresh-solve identical output, mismatched position_labels rejection.
• tests/unit/tools/test_jacobian_lens.py — unchanged behavior after the conftest extraction
  (byte-identical pass/fail set).
• tests/unit/tools/test_jacobian_lens_coordinate_patch_hooks.py — wrapper: hook-list shape mirrors
  swap_hooks, once-per-call warning naming layer/position counts, unfitted-layer rejection,
  empty-positions rejection, identical-token rejection, only-requested-positions-change, oracle
  parity with offline coordinate_patch, cache hits skip resolve across repeated hook firings,
  core errors and core warnings propagate uncaught through the hook.
• tests/integration/test_jacobian_lens.py — one cached GPT-2 test: alpha=0 no-op and
  bit-identical untouched positions through a real run_with_hooks pass.

Local gates green:

• uv run mypy . — Success: no issues found in 397 source files.
• uv run pytest tests/unit/tools/test_jacobian_lens_coordinate_patch.py tests/unit/tools/test_jacobian_lens.py tests/unit/tools/test_jacobian_lens_coordinate_patch_hooks.py — 157 passed.
• source .env && uv run pytest tests/integration/test_jacobian_lens.py -k coordinate_patch_hooks — 1 passed.

Related work

• [Proposal] Add J-space coordinate patching over sparse decompositions #1739 (closed) / feat(jacobian_lens): anchored J-space coordinate patching (offline primitive) #1741 (merged) — the offline coordinate-patching proposal and the PR that
  implemented the coordinate_patch(...) primitive this PR wraps. Strict dependency; no overlap.
• JacobianLens.swap_hooks re-reads live coordinates, causing odd/even cancellation across layer bands #1746 / Add clean-coordinate clamping for JacobianLens swaps #1747 (open) — a bug report that swap_hooks re-reads live lens coordinates at every
  hooked layer, making a multi-layer swap band involutive (odd/even cancellation), and a proposal to
  add clean-coordinate clamping for J-lens swaps. No functional conflict with this PR — different
  method, different semantics — but two things reviewers should note: (1) both areas touch
  jacobian_lens.py, the __init__.py exports, and jacobian_lens_fitting.md, so whichever lands
  second needs a mechanical rebase; and (2) coordinate_patch_hooks shares the same structural
  property — it re-decomposes the live activation at each hooked layer, so a multi-layer install
  compounds across the band. Unlike a swap, a substitute patch is not an involution, so it does not
  exhibit the exact odd/even cancellation JacobianLens.swap_hooks re-reads live coordinates, causing odd/even cancellation across layer bands #1746 describes; a clean-coordinate clamp variant of
  coordinate patching is a reasonable future PR if the same multi-layer concern proves material. This
  PR does not attempt to fix or depend on JacobianLens.swap_hooks re-reads live coordinates, causing odd/even cancellation across layer bands #1746 / Add clean-coordinate clamping for JacobianLens swaps #1747.

Out of scope (deferred to a follow-up PR)

Arbitrary multi-slot permutations, any behavioral / causal-swap benchmark (PR3), and
HookedTransformer support (PR4) remain deferred to follow-up PRs.

Checklist

• Full make test-pr run attached (unit + docstring + acceptance + integration).
• uv run build-docs run clean (new #### Dynamic coordinate-patch hooks section renders).
• Reviewer sign-off on design decisions #2 (fail-fast vs. skip-with-warning) and #3
  (caller-owned decomposition_cache shape).

[jlarson4]

Thanks for putting this together @janmenjayap! Review comments below:

[jlarson4]

Looks good, great work @janmenjayap Merging now