feat: preserve pretrained lens provenance and add filtered readouts

demos/Jacobian_Lens_Demo.ipynb

@@ -1726,12 +1726,137 @@
     "print(\"model output:\", top_j[model.cfg.n_layers - 1][-1][0])\n"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "jpp-intro",
+   "metadata": {},
+   "source": [
+    "## 7. Pretrained J++ lenses and filtered readouts\n",
+    "\n",
+    "[J++ Lens](https://github.com/safety-research/jpp_lens) (Ayonrinde & Lindsey, 2026) combines filtered Jacobian estimates into one fixed matrix per layer. We can read its published maps through the same `JacobianLens`; the LRP rules used while fitting are not needed at inference. This smaller example uses the converted Qwen3.5-0.8B artifact from [Neuronpedia](https://huggingface.co/neuronpedia/jacobian-lens), with the authors' provenance retained in `lens.metadata[\"provenance\"]`.\n",
+    "\n",
+    "The tokenizer helper excludes individually decoded tokens without a Unicode letter or digit, scanning the whole LM-head vocabulary including padding rows. `keep_token_ids` can retain scored symbols such as `+`. Special tokens with letters in their decoded spelling are retained by this policy.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "jpp-load",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "method: jpp\n",
+      "excluded vocabulary ids: 6133\n"
+     ]
+    }
+   ],
+   "source": [
+    "# NBVAL_IGNORE_OUTPUT\n",
+    "import gc\n",
+    "import torch\n",
+    "\n",
+    "from transformer_lens.model_bridge import TransformerBridge\n",
+    "from transformer_lens.tools.analysis import JacobianLens\n",
+    "from transformer_lens.utilities.tokenize_utils import readout_token_exclusions\n",
+    "\n",
+    "# Release the previous model before loading the smaller J++ example.\n",
+    "if \"lens\" in globals():\n",
+    "    lens.clear_device_cache()\n",
+    "for name in (\"result_jlens\", \"result_logitlens\", \"lens\", \"model\"):\n",
+    "    globals().pop(name, None)\n",
+    "gc.collect()\n",
+    "if torch.cuda.is_available():\n",
+    "    torch.cuda.empty_cache()\n",
+    "\n",
+    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "model = TransformerBridge.boot_transformers(\n",
+    "    \"Qwen/Qwen3.5-0.8B\",\n",
+    "    revision=\"2fc06364715b967f1860aea9cf38778875588b17\",\n",
+    "    device=device,\n",
+    "    dtype=torch.float32,\n",
+    ").eval()\n",
+    "lens = JacobianLens.from_pretrained(\n",
+    "    \"neuronpedia/jacobian-lens\",\n",
+    "    filename=\"qwen3.5-0.8b/jpp/Salesforce-wikitext/Qwen3.5-0.8B_jpp_lens.pt\",\n",
+    "    revision=\"e166deb5ad540932100ef803c8c9bcc2a5453694\",\n",
+    "    model=model,\n",
+    ")\n",
+    "assert lens.metadata[\"provenance\"][\"method\"] == \"jpp\"\n",
+    "excluded = readout_token_exclusions(model.tokenizer, vocab_size=model.W_U.shape[-1])\n",
+    "print(\"method:\", lens.metadata[\"provenance\"][\"method\"])\n",
+    "print(\"excluded vocabulary ids:\", len(excluded))\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "jpp-filter",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "layer 4\n",
+      "  raw J++: [' ,', ' \\n\\n', ' \",', ' ##', ' ,\"']\n",
+      "  filtered J++: ['<|im_end|>', ' a', ' the', ' either', ' The']\n",
+      "  filtered logit lens: ['менно', ' Nations', ' Panas', 'ң', ' geweest']\n",
+      "layer 12\n",
+      "  raw J++: [' ____', ' ___', ' ______', ' ?', ' __']\n",
+      "  filtered J++: [' nation', ' named', ' national', ' Eleven', ' Americans']\n",
+      "  filtered logit lens: ['ありますか', ' gì', 'voie', '怎样的', '其中之一']\n",
+      "layer 20\n",
+      "  raw J++: ['____', '：', '巴黎', '？', '___']\n",
+      "  filtered J++: ['巴黎', '法国', 'Paris', 'France', '哪个']\n",
+      "  filtered logit lens: ['哪一个', ' capitals', '巴黎', '哪个', 'elia']\n"
+     ]
+    }
+   ],
+   "source": [
+    "# NBVAL_IGNORE_OUTPUT\n",
+    "prompt = \"法国首都是巴黎。埃菲尔铁塔所在的城市是\"\n",
+    "readout_options = dict(\n",
+    "    layers=[4, 12, 20], positions=[-1], top_k=5, return_full_logits=True\n",
+    ")\n",
+    "raw = lens.readout(model, prompt, **readout_options)\n",
+    "filtered = lens.readout(model, prompt, exclude_token_ids=excluded, **readout_options)\n",
+    "logit_baseline = lens.readout(\n",
+    "    model, prompt, use_jacobian=False, exclude_token_ids=excluded, **readout_options\n",
+    ")\n",
+    "assert filtered.exclude_token_ids == excluded\n",
+    "assert torch.equal(raw.model_logits, filtered.model_logits)\n",
+    "assert torch.equal(raw.model_topk_indices, filtered.model_topk_indices)\n",
+    "for layer in readout_options[\"layers\"]:\n",
+    "    assert torch.equal(raw.lens_logits[layer], filtered.lens_logits[layer])\n",
+    "    assert not set(filtered.lens_topk_indices[layer].flatten().tolist()).intersection(\n",
+    "        excluded\n",
+    "    )\n",
+    "    print(f\"layer {layer}\")\n",
+    "    print(\"  raw J++:\", raw.top_tokens(model.tokenizer, k=5)[layer][0])\n",
+    "    print(\"  filtered J++:\", filtered.top_tokens(model.tokenizer, k=5)[layer][0])\n",
+    "    print(\n",
+    "        \"  filtered logit lens:\",\n",
+    "        logit_baseline.top_tokens(model.tokenizer, k=5)[layer][0],\n",
+    "    )\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "jpp-policy",
+   "metadata": {},
+   "source": [
+    "Filtering changes the ranking universe, not the raw lens logits or the model's predictions. The same exclusions are applied to the logit-lens baseline. These readouts illustrate how to use the published maps; comparing lens quality requires a fixed task evaluation rather than judging a few attractive token lists.\n"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "00df913e",
    "metadata": {},
    "source": [
-    "## 7. Fitting your own lens\n",
+    "## 8. Fitting your own lens\n",
     "\n",
     "For models without a published artifact, `JacobianLens.fit` reproduces the reference estimator on\n",
     "a raw `TransformerBridge` — one forward and `ceil(d_model / dim_batch)` backward passes per\n",

tests/unit/tools/test_jacobian_lens_filtered_readout.py

@@ -0,0 +1,141 @@
+"""Converted lens provenance and optional vocabulary exclusions."""
+
+import pytest
+import torch
+
+from tests.unit.tools.conftest import D_MODEL, D_VOCAB, N_LAYERS, _lens
+from transformer_lens.tools.analysis import JacobianLens
+
+
+def test_converted_provenance_round_trips_without_becoming_model_checks(tmp_path, toy_model):
+    provenance = {
+        "method": "jpp",
+        "source_config": {"hf_model_name": "a/mirror"},
+        "capture_point": "block_output",
+        "stream_reduce": "none",
+    }
+    path = tmp_path / "converted.pt"
+    torch.save(
+        {
+            "J": {0: torch.eye(D_MODEL)},
+            "d_model": D_MODEL,
+            "n_prompts": 96,
+            "provenance": provenance,
+            "metadata": {"note": "retained"},
+        },
+        path,
+    )
+    lens = JacobianLens.load(str(path))
+    assert lens.metadata == {"note": "retained", "provenance": provenance}
+    lens.validate_model(toy_model)
+    lens.save(str(path))
+    assert JacobianLens.load(str(path)).metadata == lens.metadata
+
+
+def test_different_imported_provenance_uses_existing_merge_guard(tmp_path):
+    lenses = []
+    for method in ("jpp", "jacobian"):
+        path = tmp_path / f"{method}.pt"
+        torch.save(
+            {
+                "J": {0: torch.eye(D_MODEL)},
+                "d_model": D_MODEL,
+                "n_prompts": 2,
+                "provenance": {"method": method},
+            },
+            path,
+        )
+        lenses.append(JacobianLens.load(str(path)))
+    with pytest.raises(ValueError, match="provenance"):
+        JacobianLens.merge(lenses)
+    merged = JacobianLens.merge([lenses[0], lenses[0]])
+    assert merged.metadata["provenance"] == {"method": "jpp"}
+
+
+@pytest.mark.parametrize("use_jacobian", [True, False])
+def test_exclusions_rank_before_topk_and_leave_raw_outputs_unchanged(toy_model, use_jacobian):
+    lens = _lens()
+    kwargs = dict(
+        layers=[0, N_LAYERS - 1],
+        positions=[-1],
+        top_k=3,
+        return_full_logits=True,
+        use_jacobian=use_jacobian,
+    )
+    baseline = lens.readout(toy_model, "a prompt", **kwargs)
+    excluded = sorted({int(ids[0, 0]) for ids in baseline.lens_topk_indices.values()})
+    filtered = lens.readout(toy_model, "a prompt", exclude_token_ids=excluded + excluded, **kwargs)
+    assert filtered.exclude_token_ids == excluded
+    assert baseline.exclude_token_ids == []
+    torch.testing.assert_close(filtered.model_logits, baseline.model_logits, rtol=0, atol=0)
+    assert torch.equal(filtered.model_topk_indices, baseline.model_topk_indices)
+    assert torch.equal(filtered.model_topk_values, baseline.model_topk_values)
+    for layer in baseline.lens_logits:
+        original = baseline.lens_logits[layer]
+        assert torch.equal(filtered.lens_logits[layer], original)
+        ranked = original.clone()
+        ranked[:, excluded] = -torch.inf
+        expected = ranked.topk(3, dim=-1)
+        assert torch.equal(filtered.lens_topk_indices[layer], expected.indices)
+        assert torch.equal(filtered.lens_topk_values[layer], original.gather(-1, expected.indices))
+
+
+def test_empty_exclusions_preserve_default_results(toy_model):
+    lens = _lens()
+    baseline = lens.readout(toy_model, "a prompt", top_k=3)
+    actual = lens.readout(toy_model, "a prompt", top_k=3, exclude_token_ids=[])
+    for layer in baseline.lens_topk_indices:
+        assert torch.equal(actual.lens_topk_indices[layer], baseline.lens_topk_indices[layer])
+        assert torch.equal(actual.lens_topk_values[layer], baseline.lens_topk_values[layer])
+
+
+def test_excluded_ids_cannot_win_negative_infinity_ties(toy_model, monkeypatch):
+    from types import SimpleNamespace
+
+    import transformer_lens.tools.analysis.jacobian_lens as module
+
+    original = toy_model.run_with_cache
+
+    def infinite_logits(*args, **kwargs):
+        logits, cache = original(*args, **kwargs)
+        return torch.full_like(logits, -torch.inf), cache
+
+    monkeypatch.setattr(toy_model, "run_with_cache", infinite_logits)
+    monkeypatch.setattr(
+        module,
+        "logit_readout",
+        lambda *args, **kwargs: SimpleNamespace(values=torch.full((1, D_VOCAB), -torch.inf)),
+    )
+    result = _lens().readout(
+        toy_model,
+        "prompt",
+        layers=[0, N_LAYERS - 1],
+        positions=[-1],
+        top_k=2,
+        exclude_token_ids=list(range(D_VOCAB - 2)),
+    )
+    for ids in result.lens_topk_indices.values():
+        assert set(ids.flatten().tolist()) == {D_VOCAB - 2, D_VOCAB - 1}
+
+
+def test_conflicting_provenance_is_not_silently_overwritten(tmp_path):
+    path = tmp_path / "conflict.pt"
+    torch.save(
+        {
+            "J": {0: torch.eye(D_MODEL)},
+            "d_model": D_MODEL,
+            "metadata": {"provenance": {"method": "j"}},
+            "provenance": {"method": "jpp"},
+        },
+        path,
+    )
+    with pytest.raises(ValueError, match="conflicting"):
+        JacobianLens.load(str(path))
+
+
+@pytest.mark.parametrize(
+    "excluded", [[-1], [D_VOCAB], [True], list(range(D_VOCAB)), list(range(D_VOCAB - 1))]
+)
+def test_invalid_or_insufficient_exclusions_raise(toy_model, excluded):
+    with pytest.raises(ValueError, match="exclude|eligible"):
+        _lens().readout(toy_model, "a prompt", top_k=2, exclude_token_ids=excluded)

tests/unit/utilities/test_readout_token_exclusions.py

@@ -0,0 +1,41 @@
+"""Unicode-aware readout filtering over the full output vocabulary."""
+
+import pytest
+
+from transformer_lens.utilities import tokenize_utils
+
+
+class Tokenizer:
+    words = ["", "\n", "!!!", "心脏", "é", "٣", "<|im_end|>", "+", " word"]
+
+    def decode(self, ids):
+        return self.words[ids[0]] if ids[0] < len(self.words) else ""
+
+
+class BatchTokenizer(Tokenizer):
+    def batch_decode(self, rows):
+        return [self.decode(row) for row in rows]
+
+
+@pytest.mark.parametrize("tokenizer", [Tokenizer(), BatchTokenizer()])
+def test_unicode_padding_and_special_tokens(tokenizer):
+    assert tokenize_utils.readout_token_exclusions(tokenizer, 11) == [0, 1, 2, 7, 9, 10]
+    assert tokenize_utils.readout_token_exclusions(tokenizer, 11, keep_token_ids=[7, 7, 3]) == [
+        0,
+        1,
+        2,
+        9,
+        10,
+    ]
+
+
+@pytest.mark.parametrize("width", [0, -1, True])
+def test_invalid_vocabulary_size(width):
+    with pytest.raises(ValueError, match="vocab_size"):
+        tokenize_utils.readout_token_exclusions(Tokenizer(), width)
+
+
+@pytest.mark.parametrize("ids", [[-1], [11], [True]])
+def test_invalid_retained_ids(ids):
+    with pytest.raises(ValueError, match="keep_token_ids"):
+        tokenize_utils.readout_token_exclusions(Tokenizer(), 11, keep_token_ids=ids)