diff --git a/docs/source/content/jacobian_lens_fitting.md b/docs/source/content/jacobian_lens_fitting.md index 86aa377b6..c0ad59141 100644 --- a/docs/source/content/jacobian_lens_fitting.md +++ b/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. diff --git a/tests/unit/tools/test_jacobian_lens.py b/tests/unit/tools/test_jacobian_lens.py index f046ad712..f096b7f02 100644 --- a/tests/unit/tools/test_jacobian_lens.py +++ b/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: diff --git a/tests/unit/tools/test_jacobian_lens_import.py b/tests/unit/tools/test_jacobian_lens_import.py new file mode 100644 index 000000000..78825a4ef --- /dev/null +++ b/tests/unit/tools/test_jacobian_lens_import.py @@ -0,0 +1,325 @@ +"""Unit tests for JacobianLens.load() checkpoint support. + +All tests use synthetic in-memory fixtures — no model, no Hub access. +""" + +import pytest +import torch + +from transformer_lens.tools.analysis.jacobian_lens import JacobianLens + +# --------------------------------------------------------------------------- +# Fixture helpers +# --------------------------------------------------------------------------- + + +def _save_artifact(path, *, n_layers=3, d_model=8, n_prompts=10, metadata=None): + """Write a minimal lens artifact (J-key format) to *path*.""" + payload = { + "J": {i: torch.randn(d_model, d_model).half() for i in range(n_layers)}, + "n_prompts": n_prompts, + "source_layers": list(range(n_layers)), + "d_model": d_model, + } + if metadata is not None: + payload["metadata"] = metadata + torch.save(payload, path) + + +def _save_checkpoint( + path, + *, + n_layers=3, + d_model=8, + n_prompts=5, + metadata=None, + flat_provenance=None, + n_prompts_override=None, +): + """Write a minimal fit-checkpoint (jacobian_sum format) to *path*. + + *n_prompts_override* lets callers inject a bad value without changing the + sums (useful for zero/negative n_prompts tests). + """ + payload = { + "jacobian_sum": {i: torch.randn(d_model, d_model) * n_prompts for i in range(n_layers)}, + "n_prompts": n_prompts if n_prompts_override is None else n_prompts_override, + "source_layers": list(range(n_layers)), + "d_model": d_model, + } + if metadata is not None: + payload["metadata"] = metadata + if flat_provenance is not None: + payload.update(flat_provenance) + torch.save(payload, path) + + +# --------------------------------------------------------------------------- +# Artifact loading (regression tests — existing behaviour must be preserved) +# --------------------------------------------------------------------------- + + +class TestLoadArtifact: + def test_basic_load(self, tmp_path): + p = str(tmp_path / "lens.pt") + _save_artifact(p, n_layers=3, d_model=8, n_prompts=10) + lens = JacobianLens.load(p) + assert lens.d_model == 8 + assert lens.n_prompts == 10 + assert lens.source_layers == [0, 1, 2] + + def test_matrices_cast_to_fp32(self, tmp_path): + p = str(tmp_path / "lens.pt") + _save_artifact(p, d_model=8) + lens = JacobianLens.load(p) + assert all(j.dtype == torch.float32 for j in lens.jacobians.values()) + + def test_metadata_preserved(self, tmp_path): + p = str(tmp_path / "lens.pt") + meta = {"model_name": "gpt2", "corpus": "wikitext"} + _save_artifact(p, metadata=meta) + lens = JacobianLens.load(p) + assert lens.metadata["model_name"] == "gpt2" + assert lens.metadata["corpus"] == "wikitext" + + def test_missing_metadata_gives_empty_dict(self, tmp_path): + p = str(tmp_path / "lens.pt") + _save_artifact(p) + lens = JacobianLens.load(p) + assert lens.metadata == {} + + def test_unknown_format_raises(self, tmp_path): + p = str(tmp_path / "garbage.pt") + torch.save({"random_key": 42}, str(p)) + with pytest.raises(ValueError, match="does not look like"): + JacobianLens.load(p) + + +# --------------------------------------------------------------------------- +# Checkpoint loading (new functionality) +# --------------------------------------------------------------------------- + + +class TestLoadCheckpoint: + def test_basic_load(self, tmp_path): + p = str(tmp_path / "ckpt.pt") + _save_checkpoint(p, n_layers=3, d_model=8, n_prompts=5) + lens = JacobianLens.load(p) + assert lens.d_model == 8 + assert lens.n_prompts == 5 + assert lens.source_layers == [0, 1, 2] + + def test_means_reconstructed_correctly(self, tmp_path): + """jacobian_sum / n_prompts must equal the stored average.""" + d_model, n_prompts = 8, 5 + sums = {i: torch.ones(d_model, d_model) * n_prompts for i in range(2)} + p = str(tmp_path / "ckpt.pt") + payload = { + "jacobian_sum": sums, + "n_prompts": n_prompts, + "source_layers": [0, 1], + "d_model": d_model, + } + torch.save(payload, p) + lens = JacobianLens.load(p) + for j in lens.jacobians.values(): + assert torch.allclose(j, torch.ones(d_model, d_model), atol=1e-6) + + def test_converted_from_key_set(self, tmp_path): + p = str(tmp_path / "ckpt.pt") + _save_checkpoint(p) + lens = JacobianLens.load(p) + assert lens.metadata.get("converted_from") == "jacobian_lens_checkpoint" + + def test_fit_reserved_key_stripped(self, tmp_path): + p = str(tmp_path / "ckpt.pt") + meta = { + "transformer_lens_fit": True, + "transformer_lens_version": "1.2.3", + "model_system": "TransformerBridge", + "corpus": "wiki", + } + _save_checkpoint(p, metadata=meta) + lens = JacobianLens.load(p) + assert "transformer_lens_fit" not in lens.metadata + assert "transformer_lens_version" not in lens.metadata + assert "model_system" not in lens.metadata + # safe keys survive + assert lens.metadata.get("corpus") == "wiki" + + def test_safe_scalar_metadata_preserved(self, tmp_path): + p = str(tmp_path / "ckpt.pt") + meta = { + "model_name": "gpt2", + "corpus": "wikitext", + "extra_int": 42, + "extra_float": 3.14, + "extra_list": [1, 2, 3], + } + _save_checkpoint(p, metadata=meta) + lens = JacobianLens.load(p) + assert lens.metadata["model_name"] == "gpt2" + assert lens.metadata["corpus"] == "wikitext" + assert lens.metadata["extra_int"] == 42 + assert lens.metadata["extra_float"] == pytest.approx(3.14) + assert lens.metadata["extra_list"] == [1, 2, 3] + + def test_tensor_valued_metadata_dropped_and_recorded(self, tmp_path): + """Tensor fields that fail _validate_metadata are logged in dropped_fields.""" + d_model, n_prompts = 8, 5 + p = str(tmp_path / "ckpt.pt") + payload = { + "jacobian_sum": {0: torch.randn(d_model, d_model) * n_prompts}, + "n_prompts": n_prompts, + "source_layers": [0], + "d_model": d_model, + "metadata": { + "model_name": "gpt2", + "embedding_stats": torch.randn(4), # tensor — will be dropped + }, + } + torch.save(payload, p) + lens = JacobianLens.load(p) + assert "embedding_stats" not in lens.metadata + dropped = lens.metadata.get("dropped_fields", []) + assert any("embedding_stats" in f for f in dropped) + # safe field survives + assert lens.metadata.get("model_name") == "gpt2" + + def test_flat_provenance_harvested(self, tmp_path): + """Checkpoints that store model_name / corpus at top level are handled.""" + p = str(tmp_path / "ckpt.pt") + _save_checkpoint( + p, flat_provenance={"model_name": "llama3", "corpus": "pile", "model_revision": "abc"} + ) + lens = JacobianLens.load(p) + assert lens.metadata.get("model_name") == "llama3" + assert lens.metadata.get("corpus") == "pile" + assert lens.metadata.get("model_revision") == "abc" + + def test_nested_metadata_overrides_flat_provenance(self, tmp_path): + """Explicit metadata dict takes precedence over flat payload keys.""" + p = str(tmp_path / "ckpt.pt") + _save_checkpoint( + p, + metadata={"model_name": "nested-name"}, + flat_provenance={"model_name": "flat-name"}, + ) + lens = JacobianLens.load(p) + assert lens.metadata.get("model_name") == "nested-name" + + def test_zero_n_prompts_raises(self, tmp_path): + p = str(tmp_path / "ckpt.pt") + _save_checkpoint(p, n_prompts_override=0) + with pytest.raises(ValueError, match="n_prompts=0"): + JacobianLens.load(p) + + def test_negative_n_prompts_raises(self, tmp_path): + p = str(tmp_path / "ckpt.pt") + _save_checkpoint(p, n_prompts_override=-1) + with pytest.raises(ValueError, match="n_prompts=-1"): + JacobianLens.load(p) + + def test_single_layer_checkpoint(self, tmp_path): + p = str(tmp_path / "ckpt.pt") + _save_checkpoint(p, n_layers=1, d_model=4) + lens = JacobianLens.load(p) + assert lens.source_layers == [0] + + def test_load_then_save_roundtrip(self, tmp_path): + """A checkpoint loaded and re-saved produces a valid artifact.""" + ckpt_path = str(tmp_path / "ckpt.pt") + art_path = str(tmp_path / "artifact.pt") + _save_checkpoint(ckpt_path, n_layers=2, d_model=8, n_prompts=3, metadata={"corpus": "wiki"}) + lens = JacobianLens.load(ckpt_path) + lens.save(art_path) + reloaded = JacobianLens.load(art_path) + assert reloaded.n_prompts == 3 + assert reloaded.source_layers == [0, 1] + assert reloaded.metadata.get("converted_from") == "jacobian_lens_checkpoint" + + +# --------------------------------------------------------------------------- +# merge() refuses to mix converted and TL-fitted lenses +# --------------------------------------------------------------------------- + + +class TestMergeProvenance: + def test_merge_two_checkpoints_succeeds(self, tmp_path): + p1, p2 = str(tmp_path / "c1.pt"), str(tmp_path / "c2.pt") + _save_checkpoint(p1, metadata={"corpus": "wiki"}) + _save_checkpoint(p2, metadata={"corpus": "wiki"}) + l1 = JacobianLens.load(p1) + l2 = JacobianLens.load(p2) + merged = JacobianLens.merge([l1, l2]) + assert merged.metadata.get("converted_from") == "jacobian_lens_checkpoint" + assert merged.n_prompts == l1.n_prompts + l2.n_prompts + + def test_merge_two_artifacts_succeeds(self, tmp_path): + shared_meta = {"model_name": "gpt2", "corpus": "wiki", "transformer_lens_fit": True} + p1, p2 = str(tmp_path / "a1.pt"), str(tmp_path / "a2.pt") + _save_artifact(p1, metadata=dict(shared_meta)) + _save_artifact(p2, metadata=dict(shared_meta)) + l1 = JacobianLens.load(p1) + l2 = JacobianLens.load(p2) + # Both have transformer_lens_fit=True — provenance matches + merged = JacobianLens.merge([l1, l2]) + assert merged.n_prompts == l1.n_prompts + l2.n_prompts + + def test_merge_checkpoint_and_artifact_raises(self, tmp_path): + """A converted checkpoint and a TL-fitted artifact must not merge.""" + p_ckpt = str(tmp_path / "ckpt.pt") + p_art = str(tmp_path / "art.pt") + _save_checkpoint(p_ckpt, metadata={"corpus": "wiki"}) + _save_artifact( + p_art, + metadata={ + "transformer_lens_fit": True, + "corpus": "wiki", + "model_name": "gpt2", + }, + ) + ckpt_lens = JacobianLens.load(p_ckpt) + art_lens = JacobianLens.load(p_art) + with pytest.raises(ValueError, match="provenance metadata"): + JacobianLens.merge([ckpt_lens, art_lens]) + + def test_merge_checkpoints_different_corpus_raises(self, tmp_path): + p1, p2 = str(tmp_path / "c1.pt"), str(tmp_path / "c2.pt") + _save_checkpoint(p1, metadata={"corpus": "wiki"}) + _save_checkpoint(p2, metadata={"corpus": "pile"}) + l1 = JacobianLens.load(p1) + l2 = JacobianLens.load(p2) + with pytest.raises(ValueError, match="provenance metadata"): + JacobianLens.merge([l1, l2]) + + def test_merge_weighted_average_correct(self, tmp_path): + """Merged matrices are prompt-count-weighted averages.""" + d_model = 4 + p1, p2 = str(tmp_path / "c1.pt"), str(tmp_path / "c2.pt") + + n1, n2 = 3, 7 + # jacobian_sum = n * mean, so here mean = ones for both + payload1 = { + "jacobian_sum": {0: torch.ones(d_model, d_model) * n1}, + "n_prompts": n1, + "d_model": d_model, + "source_layers": [0], + } + payload2 = { + "jacobian_sum": {0: torch.ones(d_model, d_model) * n2 * 2}, + "n_prompts": n2, + "d_model": d_model, + "source_layers": [0], + } + torch.save(payload1, p1) + torch.save(payload2, p2) + l1 = JacobianLens.load(p1) # mean = ones + l2 = JacobianLens.load(p2) # mean = 2*ones + merged = JacobianLens.merge([l1, l2]) + expected_mean = (n1 * 1.0 + n2 * 2.0) / (n1 + n2) + assert torch.allclose( + merged.jacobians[0], + torch.full((d_model, d_model), expected_mean), + atol=1e-5, + ) diff --git a/transformer_lens/tools/analysis/jacobian_lens.py b/transformer_lens/tools/analysis/jacobian_lens.py index f93f21d79..e333ffd1d 100644 --- a/transformer_lens/tools/analysis/jacobian_lens.py +++ b/transformer_lens/tools/analysis/jacobian_lens.py @@ -121,6 +121,29 @@ def _resolve_registry_entry(name_or_path: str) -> Optional[Tuple[str, str]]: _SWAP_WARN_COSINE = 0.99 _SWAP_ERROR_COSINE = 0.999 +# Keys written by fit() that must not appear in converted-lens metadata so that +# merge() can refuse to mix TL-fitted lenses with externally converted ones. +# Note: "target_layer" is intentionally NOT listed here — it must survive +# conversion so that validate_model() can detect and refuse checkpoints that +# were fitted against a non-final target layer. +_FIT_RESERVED_KEYS: frozenset = frozenset( + { + "transformer_lens_fit", + "transformer_lens_version", + "model_system", + "processing", + "hook_convention", + "fit_dtype", + "dim_batch", + "max_seq_len", + "skip_first_positions", + } +) + +# Top-level payload keys that some checkpoint writers store as flat provenance +# (rather than nested under a "metadata" key). +_CHECKPOINT_FLAT_PROVENANCE: frozenset = frozenset({"model_name", "model_revision", "corpus"}) + @dataclass class JacobianLensReadout: @@ -266,28 +289,136 @@ def save(self, path: str, *, dtype: torch.dtype = torch.float16) -> None: @classmethod def load(cls, path: str) -> "JacobianLens": - """Load a lens artifact saved by this class or the reference package. + """Load a lens artifact or fit checkpoint saved in a supported schema. + + Two file schemas are accepted: + + **Artifact** (the reference format, written by :meth:`save` or the + Anthropic reference package): must contain a ``J`` key mapping layer + indices to transport matrices, plus ``n_prompts``, ``d_model``, and an + optional ``metadata`` dict. + + **Fit checkpoint** (running-sum format, written by the reference + implementation's ``write_checkpoint()`` during fitting): must contain + a ``jacobian_sum`` key mapping layer indices to *running-sum* matrices + (i.e. the sum over prompts, not yet divided by the prompt count), plus + ``n_done``. ``d_model`` is inferred from the first matrix's shape; + no explicit ``d_model`` key is required or expected. The per-layer + means are reconstructed on load. A + ``converted_from: "jacobian_lens_checkpoint"`` key is added to + metadata so :meth:`merge` refuses to silently combine checkpoints with + natively TL-fitted lenses. Fit-reserved provenance keys + (``transformer_lens_fit``, etc.) are stripped; scalar fields that can + be serialised under ``weights_only=True`` are preserved. + Tensor-valued metadata fields that would fail :func:`_validate_metadata` + are recorded by name in a ``dropped_fields`` list. + + Fit checkpoint schema (reference ``write_checkpoint()`` format) + --------------------------------------------------------------- + The reference implementation writes exactly six top-level keys; all + other keys in the payload are ignored:: + + { + "jacobian_sum": {: , ...}, + "n_done": , # prompts accumulated into jacobian_sum + "next_idx": , # next prompt index (informational) + "source_layers": [, ...], # documented layer indices (informational) + "target_layer": , # target layer — harvested into metadata + "skip_first": , # leading positions skipped (informational) + # optional flat provenance accepted from alternative checkpoint writers: + "model_name": , + "model_revision": , + "corpus": , + # optional nested provenance accepted from alternative writers: + "metadata": {: , ...}, + } Args: - path: Path to the ``.pt`` artifact. + path: Path to the ``.pt`` file. Raises: - ValueError: If the file lacks the ``J`` key (e.g. a fit checkpoint - rather than a saved lens). + ValueError: If the file lacks both a ``J`` key (artifact) and a + ``jacobian_sum`` key (checkpoint), or if a checkpoint records a + non-positive ``n_prompts``. """ payload = torch.load(path, map_location="cpu", weights_only=True) - if "J" not in payload: - raise ValueError( - f"{path} does not look like a Jacobian lens artifact (no 'J' key). " - "Fit checkpoints and other formats are not supported." + if "J" in payload: + return cls( + {int(layer): matrix for layer, matrix in payload["J"].items()}, + n_prompts=int(payload.get("n_prompts", 0)), + d_model=int(payload["d_model"]), + metadata=payload.get("metadata"), ) - return cls( - {int(layer): matrix for layer, matrix in payload["J"].items()}, - n_prompts=int(payload.get("n_prompts", 0)), - d_model=int(payload["d_model"]), - metadata=payload.get("metadata"), + if "jacobian_sum" in payload: + return cls._from_checkpoint_payload(path, payload) + raise ValueError( + f"{path} does not look like a Jacobian lens artifact or fit checkpoint. " + "Expected a 'J' key (artifact) or 'jacobian_sum' key (fit checkpoint). " + "See JacobianLens.load() for the supported file schemas." ) + @classmethod + def _from_checkpoint_payload(cls, path: str, payload: Dict[str, Any]) -> "JacobianLens": + """Reconstruct a JacobianLens from a fit-checkpoint payload. + + Divides the running Jacobian sums by ``n_prompts``, strips fit-reserved + provenance keys, harvests safe scalar metadata, records dropped tensor + fields, and marks the result as converted so :meth:`merge` refuses to + mix it with natively TL-fitted lenses. + """ + n_prompts = int(payload.get("n_done", payload.get("n_prompts", 0))) + if n_prompts <= 0: + raise ValueError( + f"{path} is a fit checkpoint with n_prompts={n_prompts}; " + "a positive prompt count is required to reconstruct the Jacobian mean." + ) + if not payload["jacobian_sum"]: + raise ValueError( + f"{path} is a fit checkpoint with an empty jacobian_sum; " + "at least one layer matrix is required to reconstruct d_model." + ) + first_matrix = next(iter(payload["jacobian_sum"].values())) + d_model = first_matrix.shape[0] + jacobians = { + int(layer): matrix.float() / n_prompts + for layer, matrix in payload["jacobian_sum"].items() + } + + # Collect raw provenance: first from nested "metadata", then supplement + # with flat top-level keys that some checkpoint writers place directly + # in the payload (model_name, model_revision, corpus). + raw_meta: Dict[str, Any] = dict(payload.get("metadata") or {}) + for key in _CHECKPOINT_FLAT_PROVENANCE: + if key in payload and key not in raw_meta: + raw_meta[key] = payload[key] + # target_layer lives at the top level in the reference checkpoint format; + # harvest it into metadata so validate_model() can check the fitting target. + if "target_layer" in payload and "target_layer" not in raw_meta: + raw_meta["target_layer"] = payload["target_layer"] + + # Build clean metadata: drop fit-reserved keys, record tensor-valued + # fields that _validate_metadata would reject (they cannot survive a + # weights_only=True reload), and keep everything else that validates. + dropped_fields: List[str] = [] + clean_meta: Dict[str, Any] = {} + for key, value in raw_meta.items(): + if key in _FIT_RESERVED_KEYS: + continue + if isinstance(value, torch.Tensor): + dropped_fields.append(f"{key}: shape={tuple(value.shape)} dtype={value.dtype}") + continue + try: + _validate_metadata({key: value}) + clean_meta[key] = value + except ValueError: + dropped_fields.append(key) + + clean_meta["converted_from"] = "jacobian_lens_checkpoint" + if dropped_fields: + clean_meta["dropped_fields"] = dropped_fields + + return cls(jacobians, n_prompts=n_prompts, d_model=d_model, metadata=clean_meta) + @classmethod def from_pretrained( cls,