[aidxhxr]

Description

JacobianLens.load() reads the matrices of the published J++ lenses but drops their top-level provenance dict, so the released Qwen3.5-0.8B file on neuronpedia/jacobian-lens loads with metadata == {}. The authors' own file format from safety-research/jpp_lens (parameters.jacobians + config) is rejected outright. This is part 1 of #1991 (loader and provenance); the filtered-readout API from part 2 can follow separately.

What changes:

• load() keeps a recorded provenance dict verbatim under metadata["provenance"] and normalises the values it records (model name, revision when present, method, lrp_mode, target layer or relative end-transport layer, capture point, stream reduction) onto the metadata keys validate_model() already reads. Explicit metadata keys win. Nothing is invented: a file with no recorded revision loads without model_revision. Leaves that weights_only=True could not reload are dropped and named in dropped_fields.
• load() accepts the J++ Lens package schema (parameters.jacobians, source_layers, config) and its {"J", "config"} variant, reading d_model and num_prompts_trained_on from the config, keeping the config under metadata["provenance"]["source_config"], and marking the lens converted_from: "jpp_lens_artifact". Only lens_type == "jacobian" files load.
• validate_model() resolves relative_end_transport_layer against the model and checks a recorded capture_point (block_output) and stream_reduce (none) against the single-stream block-output contract.
• merge() refuses imported J++ maps (method == "jpp" or the sentinel). They combine clustered expert Jacobians with learned weights, so they are not prompt means and carry no additive statistics.
• Plain artifacts, TL-fitted artifacts and fit checkpoints load exactly as before.

Checked against the real pinned artifact (neuronpedia/jacobian-lens @ e166deb5, qwen3.5-0.8b/jpp/.../Qwen3.5-0.8B_jpp_lens.pt, SHA-256 39bf7690...): metadata is now populated, save/load round trips exactly, and merge() refuses it.

Fixes part 1 of #1991

Type of change

• New feature (non-breaking change which adds functionality)
• This change requires a documentation update

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

[arrdel]

Checked this against safety-research/jpp_lens (workspace_lens/config.py and workspace_lens/lenses/base_lens.py). The schema split (J + provenance, J + config, and parameters.jacobians + config), the merge() guard on non-additive maps, and the relative_end_transport_layer / capture_point / stream_reduce checks in validate_model() all match the upstream package.

Two edge cases against jpp_lens's on-disk format worth tightening:

1. LensConfig defaults d_model to 0, not None. In workspace_lens/config.py (LensConfig), the untrained default is d_model: int = 0, so config.to_dict() writes "d_model": 0 rather than omitting the key. In both _from_artifact_payload and _from_jpp_payload, the fallback check is if d_model is None: d_model = _matrix_width(path, jacobians). When a file carries "d_model": 0, 0 is not None skips _matrix_width and passes d_model=0 into cls(...), which then raises jacobians[0] has shape (d, d), expected (0, 0). Checking if not d_model: (or if d_model is None or int(d_model) <= 0:) in both loaders handles both an omitted key and the dataclass default 0.

2. source_layers consistency check in _from_jpp_payload. BaseLens.save() always writes top-level "source_layers": self.source_layers alongside "parameters", and LensParameters.__post_init__ in base_lens.py enforces sorted(per_layer) == sorted(source_layers). _from_jpp_payload's docstring notes that the file holds the source_layers list, but the loader ignores payload.get("source_layers"). Checking that sorted(jacobians) agrees with payload["source_layers"] when present (and raising ValueError on mismatch, as BaseLens.load does) prevents silently loading a truncated or mismatched layer dict.

[aidxhxr]

Thanks, both are right and fixed in 820159d. Both loaders now treat d_model: 0 (the LensConfig default) as unset and fall back to the matrix width, and _from_jpp_payload raises ValueError when a top-level source_layers disagrees with the layers in parameters.jacobians, as BaseLens.load does. Tests added for both.

[jlarson4]

Closing as superseded by #2014