support loading fit checkpoints in JacobianLens.load()

docs/source/content/jacobian_lens_fitting.md

@@ -232,3 +232,93 @@ lens = JacobianLens.from_pretrained(
 To propose a short-name entry in TransformerLens, open a pull request that adds the
 published file to `transformer_lens/tools/analysis/jacobian_lens_registry.json` and
 include the fitting provenance and validation results.
+
+## Importing an existing lens
+
+`JacobianLens.load()` accepts two file schemas: the standard artifact format (four
+official keys — `J`, `n_prompts`, `source_layers`, `d_model`) written by
+`JacobianLens.save()` and compatible with the Anthropic reference package, and the
+fit-checkpoint format described below.
+
+### Fit checkpoint format
+
+A fit checkpoint holds the running Jacobian sums from an interrupted or staged fit,
+before the final per-prompt average is applied. `JacobianLens.load()` detects a
+checkpoint by the presence of a `jacobian_sum` key and reconstructs the per-layer
+means automatically:
+
+```python
+lens = JacobianLens.load("path/to/checkpoint.pt")
+```
+
+The reference implementation (`anthropics/jacobian-lens`) writes exactly six
+top-level keys. `load()` also accepts optional flat-provenance keys from alternative
+writers. The full set of recognised keys is:
+
+| Key | Type | Required | Description |
+|---|---|---|---|
+| `jacobian_sum` | `dict[int, Tensor[d, d]]` | yes | Running sum of per-prompt Jacobians (float32) |
+| `n_done` | `int` | yes | Number of prompts accumulated into the sums (must be > 0) |
+| `next_idx` | `int` | no | Index of the next prompt to process (informational; not used by `load()`) |
+| `source_layers` | `list[int]` | no | Documented layer indices (informational only) |
+| `target_layer` | `int` | no | Target layer fitted against — harvested into `metadata` so `validate_model()` can check it |
+| `skip_first` | `int` | no | Leading positions excluded from the source average (informational; not used by `load()`) |
+| `model_name` | `str` | no | Flat provenance shortcut (harvested into `metadata`) |
+| `model_revision` | `str` | no | Flat provenance shortcut (harvested into `metadata`) |
+| `corpus` | `str` | no | Flat provenance shortcut (harvested into `metadata`) |
+| `metadata` | `dict[str, ...]` | no | Nested provenance from alternative writers (scalars, lists, string-keyed dicts) |
+
+`d_model` is **not** read from the payload; it is inferred from the shape of the
+first matrix in `jacobian_sum`.
+
+`load()` always sets `converted_from: "jacobian_lens_checkpoint"` in the resulting
+lens's metadata. This sentinel prevents `merge()` from silently combining a
+converted checkpoint with a natively TL-fitted lens, since the two provenance
+dictionaries will differ. To merge shards that were all loaded from checkpoints, they
+must otherwise share identical provenance (same `model_name`, `corpus`, etc.):
+
+```python
+shards = [JacobianLens.load(str(p)) for p in checkpoint_paths]
+# All shards get converted_from="jacobian_lens_checkpoint";
+# merge() requires identical provenance apart from n_prompts.
+merged = JacobianLens.merge(shards)
+merged.save("merged_lens.pt")
+```
+
+### Metadata handling
+
+Three things are guaranteed when loading a checkpoint:
+
+**Fit-reserved keys are stripped.** Keys written by `JacobianLens.fit()` to record
+TransformerLens internals (`transformer_lens_fit`, `transformer_lens_version`,
+`model_system`, `hook_convention`, etc.) are not carried over, since they describe
+the *native* fitting pipeline and would be wrong for an imported lens.
+`target_layer` is a deliberate exception: it is preserved in the resulting metadata
+so `validate_model()` can detect and reject checkpoints fitted against a non-final
+target layer.
+
+**Tensor-valued fields are dropped.** `JacobianLens.save()` uses
+`weights_only=True` for reload safety, which restricts metadata to plain Python
+scalars, lists, and string-keyed dicts. Tensor-valued fields from the source
+checkpoint would silently corrupt a future load, so they are dropped and their names
+are recorded in a `dropped_fields` list in the lens's metadata:
+
+```python
+lens = JacobianLens.load("checkpoint_with_tensors.pt")
+if "dropped_fields" in lens.metadata:
+    print("fields not carried over:", lens.metadata["dropped_fields"])
+```
+
+**Dtype is not preserved.** `JacobianLens.__init__` casts all matrices to float32
+and `save()` defaults to float16, so the source checkpoint's storage dtype cannot
+be preserved end-to-end. Value precision is preserved within those two conversions.
+
+### Note on tuned-lens
+
+Tuned-lens checkpoints are not currently supported. Tuned-lens translators are
+affine (weight + bias), but the Jacobian artifact format has no bias slot to
+receive the translation component; importing would therefore silently drop the
+bias and produce wrong-but-plausible readouts. Layer indexing also differs
+(input-to-layer-ℓ vs. output-of-block-ℓ), which would yield misaligned
+readouts if not corrected. These issues are deferred; support can be added in a
+follow-up once a lossless mapping is established.

tests/unit/tools/test_jacobian_lens.py

@@ -501,13 +501,137 @@ def test_load_official_format_without_metadata(tmp_path: Any) -> None:
     assert lens.jacobians[0].dtype == torch.float32
 
 
-def test_load_rejects_fit_checkpoints(tmp_path: Any) -> None:
+def test_load_checkpoint_with_zero_n_done_raises(tmp_path: Any) -> None:
+    """A checkpoint with n_done=0 cannot reconstruct the Jacobian mean."""
     path = str(tmp_path / "ckpt.pt")
-    torch.save({"jacobian_sum": {}, "n_done": 3}, path)
-    with pytest.raises(ValueError, match="no 'J' key"):
+    torch.save(
+        {
+            "jacobian_sum": {0: torch.zeros(D_MODEL, D_MODEL)},
+            "n_done": 0,
+            "next_idx": 0,
+            "source_layers": [0],
+            "target_layer": N_LAYERS - 1,
+            "skip_first": 0,
+        },
+        path,
+    )
+    with pytest.raises(ValueError, match="n_prompts=0"):
+        JacobianLens.load(path)
+
+
+def test_load_checkpoint_with_empty_jacobian_sum_raises(tmp_path: Any) -> None:
+    """A checkpoint with an empty jacobian_sum cannot derive d_model."""
+    path = str(tmp_path / "ckpt.pt")
+    torch.save(
+        {
+            "jacobian_sum": {},
+            "n_done": 3,
+            "next_idx": 3,
+            "source_layers": [],
+            "target_layer": N_LAYERS - 1,
+            "skip_first": 0,
+        },
+        path,
+    )
+    with pytest.raises(ValueError, match="empty jacobian_sum"):
         JacobianLens.load(path)
 
 
+def test_load_checkpoint_mirrors_fit_payload_schema(tmp_path: Any) -> None:
+    """Fixture uses the verbatim 6-key schema that write_checkpoint() produces.
+
+    The reference implementation (anthropics/jacobian-lens) writes exactly:
+    jacobian_sum, n_done, next_idx, source_layers, target_layer, skip_first —
+    no ``d_model`` key, no nested ``metadata`` dict.
+
+    After load():
+    - n_done drives the prompt count
+    - d_model is inferred from the first jacobian_sum matrix's shape
+    - target_layer is harvested from the top-level payload into metadata
+    - converted_from sentinel is set so merge() refuses to mix with fitted lenses
+    """
+    path = str(tmp_path / "fit_checkpoint.pt")
+    torch.save(
+        {
+            "jacobian_sum": {0: torch.ones(D_MODEL, D_MODEL) * 5.0},
+            "n_done": 5,
+            "next_idx": 5,
+            "source_layers": [0],
+            "target_layer": N_LAYERS - 1,
+            "skip_first": 0,
+        },
+        path,
+    )
+    lens = JacobianLens.load(path)
+
+    # jacobian_sum / n_done = ones*5 / 5 = ones
+    assert torch.allclose(lens.jacobians[0], torch.ones(D_MODEL, D_MODEL))
+    assert lens.n_prompts == 5
+    # d_model derived from matrix shape, not a payload key
+    assert lens.d_model == D_MODEL
+
+    # target_layer harvested from top-level payload into metadata
+    assert lens.metadata["target_layer"] == N_LAYERS - 1
+
+    # converted_from sentinel must be set so merge() refuses to mix with fitted lenses
+    assert lens.metadata["converted_from"] == "jacobian_lens_checkpoint"
+
+
+def test_load_checkpoint_harvests_flat_provenance_and_strips_fit_keys(
+    tmp_path: Any,
+) -> None:
+    """Flat provenance and nested fit-reserved keys are handled correctly.
+
+    Some checkpoint writers (e.g. alternative TL-based fitters) may include flat
+    provenance keys (model_name, corpus, model_revision) at the top level and/or
+    a nested metadata dict with fit-reserved fields. Verify the loader strips
+    fit-reserved keys and harvests safe provenance regardless of nesting.
+    """
+    path = str(tmp_path / "tl_checkpoint.pt")
+    torch.save(
+        {
+            "jacobian_sum": {0: torch.ones(D_MODEL, D_MODEL) * 5.0},
+            "n_done": 5,
+            "next_idx": 5,
+            "source_layers": [0],
+            "target_layer": N_LAYERS - 1,
+            "skip_first": 0,
+            # flat provenance at top level
+            "model_name": "toy-bridge",
+            "model_revision": "abc123def456",
+            "corpus": CORPUS,
+            # nested metadata with fit-reserved fields as an alternative writer might add
+            "metadata": {
+                "transformer_lens_fit": True,
+                "transformer_lens_version": "3.7.0",
+                "hook_convention": "blocks.{layer}.hook_out",
+                "fit_dtype": "float32",
+                "dim_batch": 8,
+                "max_seq_len": 128,
+                "skip_first_positions": 16,
+            },
+        },
+        path,
+    )
+    lens = JacobianLens.load(path)
+
+    # flat provenance harvested into metadata
+    assert lens.metadata["model_name"] == "toy-bridge"
+    assert lens.metadata["corpus"] == CORPUS
+
+    # fit-reserved internal keys must be stripped
+    assert "transformer_lens_fit" not in lens.metadata
+    assert "transformer_lens_version" not in lens.metadata
+    assert "hook_convention" not in lens.metadata
+    assert "fit_dtype" not in lens.metadata
+
+    # target_layer must be preserved (NOT stripped) so validate_model can check it
+    assert lens.metadata["target_layer"] == N_LAYERS - 1
+
+    # converted_from sentinel must be set
+    assert lens.metadata["converted_from"] == "jacobian_lens_checkpoint"
+
+
 def test_from_pretrained_uses_retry_and_forwards_hub_arguments(
     monkeypatch: pytest.MonkeyPatch, tmp_path: Any
 ) -> None: