diff --git a/README.md b/README.md
index 1963cfd..3462bdd 100644
--- a/README.md
+++ b/README.md
@@ -92,62 +92,65 @@ python -m scripts.workspace_geometry \
 ## What the code looks like
 
 ```python
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from jspace.model_adapter import load_model_and_tokenizer, layer_indices
+import torch
+
+from jspace.model_adapter import (
+    cache_residuals,
+    get_unembedding_matrix,
+    layer_indices,
+    load_model,
+    normalize_fn,
+)
 from jspace.jacobian_lens import train_jacobian_lens
 from jspace.readout import lens_readout
-from jspace.decomposition import decompose_jspace, jspace_occupancy
+from jspace.decomposition import decompose_jspace
 from jspace.interventions import coordinate_swap, apply_intervention
 
-model, tokenizer = load_model_and_tokenizer("gpt2")
-text = "The capital of France is"
-inputs = tokenizer(text, return_tensors="pt")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model, tokenizer = load_model("gpt2", device)
 
 # Train J_l matrices for every layer (cached on disk)
+texts = ["The capital of France is", "The capital of Japan is"]
+enc = tokenizer(texts, return_tensors="pt", padding=True)
 J = train_jacobian_lens(
     model,
-    tokenizer,
-    corpus_inputs=inputs,
+    enc["input_ids"],
+    attention_mask=enc["attention_mask"],
     target_layer=layer_indices(model)[-2],  # penultimate layer target
 )
 
+W_U = get_unembedding_matrix(model)
+norm = normalize_fn(model)
+
 # Read out layer 8 as if it were the final logits
-readout_probs = lens_readout(
-    model,
-    hidden_state=...,                       # h_8 from a forward pass
-    J_l=J[8],
-    tokenizer=tokenizer,
+residuals = cache_residuals(
+    model, enc["input_ids"], layers=[8], attention_mask=enc["attention_mask"]
 )
-print(tokenizer.decode(readout_probs.argmax(dim=-1)))
+h_8 = residuals[8][0, -1]  # last position of the first sequence
+top_ids, top_probs = lens_readout(h_8, J[8], W_U, norm)
+print(tokenizer.decode(top_ids.tolist()))
 
 # Decompose a hidden state into sparse J-space coefficients
-# Pass the pre-built dictionary V = W_U @ J_l ...
-coeffs, h_J, h_perp = decompose_jspace(
-    hidden_state=...,
-    V=V,
-    k=10,
-    non_negative=True,
-)
+# Pass the pre-built dictionary V = W_U @ J_l (rows are the J-lens vectors) ...
+V = W_U @ J[8]
+coeffs, h_J, h_perp = decompose_jspace(h_8, V=V, k=10)
 
 # ...or pass J_l and W_U and let the function build V for you.
-coeffs, h_J, h_perp = decompose_jspace(
-    hidden_state=...,
-    J_l=J[8],
-    W_U=get_unembedding_matrix(model),
-    k=10,
-    non_negative=True,
-)
+coeffs, h_J, h_perp = decompose_jspace(h_8, J_l=J[8], W_U=W_U, k=10)
 
 # Swap the J-space coordinates of two tokens and run the model with the edit
-edited = apply_intervention(
+paris_id = tokenizer.encode(" Paris", add_special_tokens=False)[0]
+london_id = tokenizer.encode(" London", add_special_tokens=False)[0]
+
+def swap(h: torch.Tensor, _pos: int) -> torch.Tensor:
+    return coordinate_swap(h, paris_id, london_id, V)
+
+edited_logits = apply_intervention(
     model,
-    inputs,
-    intervention=coordinate_swap,
+    swap,
     layer_band=(6, 10),
-    J=J,
-    tokenizer=tokenizer,
-    source_pos=2,
-    target_pos=4,
+    input_ids=enc["input_ids"],
+    attention_mask=enc["attention_mask"],
 )
 ```
 
diff --git a/demo.ipynb b/demo.ipynb
index e07ac3f..b540ee4 100644
--- a/demo.ipynb
+++ b/demo.ipynb
@@ -52,7 +52,8 @@
     "    \"Once upon a time, a curious robot decided to\",\n",
     "] * 4\n",
     "enc = tokenizer(prompts, return_tensors=\"pt\", padding=True, truncation=True, max_length=32)\n",
-    "corpus_ids = enc[\"input_ids\"]"
+    "corpus_ids = enc[\"input_ids\"]\n",
+    "corpus_mask = enc[\"attention_mask\"]"
    ]
   },
   {
@@ -65,6 +66,7 @@
     "J = train_jacobian_lens(\n",
     "    model,\n",
     "    corpus_ids,\n",
+    "    attention_mask=corpus_mask,\n",
     "    target_layer=layers[-2],\n",
     "    dtype=torch.float32,\n",
     "    max_positions=32,\n",
@@ -122,7 +124,7 @@
    "outputs": [],
    "source": [
     "metrics = compute_discovery_metrics(\n",
-    "    model, tokenizer, J, corpus_ids, W_U, normalize_fn(model)\n",
+    "    model, tokenizer, J, corpus_ids, W_U, normalize_fn(model), attention_mask=corpus_mask\n",
     ")\n",
     "start, end = infer_workspace_boundaries(metrics)\n",
     "get_console().print(metrics_table(layers, metrics, (start, end)))\n"
diff --git a/jspace/decomposition.py b/jspace/decomposition.py
index 8a8b3cf..6404900 100644
--- a/jspace/decomposition.py
+++ b/jspace/decomposition.py
@@ -101,6 +101,9 @@ def jspace_occupancy(
     prev_error = float(np.linalg.norm(residual))
     selected: list[int] = []
     for k in range(max_k):
+        if prev_error == 0.0:
+            # Residual exhausted: no (k+1)-th vector can improve anything.
+            return k
         corr = V_np @ residual
         # Zero out (not -inf) already-selected atoms: the |corr| argmax below
         # would otherwise turn -inf into the winning entry.
diff --git a/jspace/discovery.py b/jspace/discovery.py
index 906c027..b9da48c 100644
--- a/jspace/discovery.py
+++ b/jspace/discovery.py
@@ -8,6 +8,8 @@
 from scipy.stats import kurtosis as scipy_kurtosis
 from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
 
+from jspace import JSpaceError
+
 # Thresholds used to locate the workspace / motor-onset boundaries.
 ACCURACY_MOTOR_THRESHOLD = 0.8
 KURTOSIS_STD_MULTIPLIER = 0.5
@@ -40,6 +42,7 @@ def compute_discovery_metrics(
     norm_fn: Callable[[torch.Tensor], torch.Tensor],
     layers: list[int] | None = None,
     probe_ids: torch.Tensor | None = None,
+    attention_mask: torch.Tensor | None = None,
 ) -> dict[str, np.ndarray]:
     """Compute CKA, kurtosis, accuracy, and autocorrelation by layer.
 
@@ -47,9 +50,21 @@ def compute_discovery_metrics(
     J-Lens was trained only up to a target layer.  If `probe_ids` is provided,
     CKA is computed over only those token rows (matching the article's shared
     probe-token geometry); otherwise the full vocabulary is used.
+
+    Pass the tokenizer's ``attention_mask`` alongside ``corpus`` so padded
+    positions are excluded from the metrics; without it, pad targets are
+    filtered via ``tokenizer.pad_token_id`` where defined. When a mask is
+    provided it is forwarded to the model and all metrics are computed over
+    unmasked positions only.
     """
     from jspace.model_adapter import layer_indices
 
+    if attention_mask is not None and attention_mask.shape != corpus.shape:
+        raise JSpaceError(
+            f"attention_mask shape {tuple(attention_mask.shape)} does not match "
+            f"corpus shape {tuple(corpus.shape)}"
+        )
+
     device = next(model.parameters()).device
     if layers is None:
         layers = layer_indices(model)
@@ -74,19 +89,33 @@ def compute_discovery_metrics(
 
     model.eval()
     with torch.no_grad():
-        corpus_batch = corpus[: min(corpus.shape[0], 32)].to(device)
-        full_outputs = model(corpus_batch, output_hidden_states=True, return_dict=True)
+        n_rows = min(corpus.shape[0], 32)
+        corpus_batch = corpus[:n_rows].to(device)
+        mask_batch = attention_mask[:n_rows].to(device) if attention_mask is not None else None
+        full_outputs = model(
+            corpus_batch,
+            attention_mask=mask_batch,
+            output_hidden_states=True,
+            return_dict=True,
+        )
         hidden_states = full_outputs.hidden_states
         targets = corpus_batch[:, 1:].contiguous()
 
     for idx, layer in enumerate(layers):
         h_last = hidden_states[layer + 1][:, :-1, :].reshape(-1, hidden_states[layer + 1].shape[-1])
         flat_targets = targets.reshape(-1)
-        valid = (flat_targets != tokenizer.pad_token_id).to(device)
+        if mask_batch is not None:
+            valid = mask_batch[:, 1:].reshape(-1).to(torch.float32)
+        else:
+            valid = (flat_targets != tokenizer.pad_token_id).to(device)
 
         logits = F.linear(norm_fn(h_last), W_U.to(device, torch.float32))
         probs = F.softmax(logits, dim=-1)
-        kurt[idx] = _excess_kurtosis(probs)
+        if mask_batch is not None:
+            keep = valid.bool()
+            kurt[idx] = _excess_kurtosis(probs[keep]) if bool(keep.any()) else float("nan")
+        else:
+            kurt[idx] = _excess_kurtosis(probs)
 
         preds = logits.argmax(dim=-1)
         correct = ((preds == flat_targets) & valid.bool()).float().sum()
@@ -94,8 +123,14 @@ def compute_discovery_metrics(
 
         top1 = preds.reshape(corpus_batch.shape[0], -1)
         shifted = torch.roll(top1, shifts=-1, dims=1)
-        matches = (top1[:, :-1] == shifted[:, :-1]).float()
-        autocorr[idx] = float(matches.mean())
+        matches = top1[:, :-1] == shifted[:, :-1]
+        if mask_batch is not None:
+            pair_valid = (mask_batch[:, 1:-1] & mask_batch[:, 2:]).bool()
+            autocorr[idx] = float(
+                (matches & pair_valid).float().sum() / (pair_valid.float().sum() + 1e-9)
+            )
+        else:
+            autocorr[idx] = float(matches.float().mean())
 
     return {"cka_block": cka, "kurtosis": kurt, "accuracy": acc, "autocorr": autocorr}
 
diff --git a/jspace/jacobian_lens.py b/jspace/jacobian_lens.py
index 63f20cf..bc01d5c 100644
--- a/jspace/jacobian_lens.py
+++ b/jspace/jacobian_lens.py
@@ -10,7 +10,7 @@
 
 from jspace import JSpaceError
 from jspace.model_adapter import _base_model, _layer_container, _norm_module, layer_indices
-from jspace.utils import get_position_ids, lens_cache_exists, load_lens_layer, save_lens_layer
+from jspace.utils import lens_cache_exists, load_lens_layer, save_lens_layer
 from jspace.viz import jl_track
 
 
@@ -94,6 +94,11 @@ def hook(module, inp, out):
     return handles
 
 
+def _position_ids_from_mask(attention_mask: torch.Tensor) -> torch.Tensor:
+    """Positions derived from an attention mask, clamped so left padding maps to 0."""
+    return (attention_mask.cumsum(dim=-1) - 1).clamp_min(0)
+
+
 def _capture_h_l(
     model: nn.Module,
     input_ids: torch.Tensor,
@@ -101,7 +106,35 @@ def _capture_h_l(
     layer_idx: int,
 ) -> torch.Tensor:
     """Capture the residual stream after layer_idx (post-layer residual)."""
-    position_ids = get_position_ids(attention_mask)
+    position_ids = _position_ids_from_mask(attention_mask)
+    num_layers = len(layer_indices(model))
+    if layer_idx + 1 == num_layers:
+        # transformers >= 5.x ties hidden_states[-1] to the post-norm
+        # last_hidden_state, so capture the pre-norm residual via the norm input.
+        final_norm = _norm_module(model)
+        if final_norm is None:
+            raise JSpaceError("Could not locate final normalization module")
+        norm_input: list[torch.Tensor | None] = [None]
+
+        def capture_norm_input(module, input_tuple):
+            norm_input[0] = input_tuple[0]
+            return input_tuple
+
+        handle = final_norm.register_forward_pre_hook(capture_norm_input)
+        try:
+            with torch.no_grad():
+                model(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    return_dict=True,
+                )
+        finally:
+            handle.remove()
+        if norm_input[0] is None:
+            raise JSpaceError("Failed to capture pre-final-norm residual")
+        return norm_input[0]
+
     with torch.no_grad():
         out = model(
             input_ids=input_ids,
@@ -131,7 +164,7 @@ def _run_from_layer(
     """
     base = _base_model(model)
     num_layers = len(layer_indices(model))
-    position_ids = get_position_ids(attention_mask)
+    position_ids = _position_ids_from_mask(attention_mask)
 
     if target_layer is None:
         capture_layer = num_layers - 1
@@ -305,7 +338,6 @@ def train_jacobian_lens(
             batches = range(0, corpus.shape[0], batch_size)
             for start in jl_track(batches, "Training J-Lens"):
                 batch_ids = corpus[start : start + batch_size]
-                B, T = batch_ids.shape
                 if attention_mask is not None:
                     batch_mask = attention_mask[start : start + batch_size]
                 else:
@@ -316,6 +348,14 @@ def train_jacobian_lens(
                 device = next(model.parameters()).device
                 batch_ids = batch_ids.to(device)
                 batch_mask = batch_mask.to(device)
+                # Rows with no valid positions contribute nothing; dropping them
+                # keeps their zeros out of the J_l average and its divisor.
+                keep = batch_mask.sum(dim=1) > 0
+                batch_ids = batch_ids[keep]
+                batch_mask = batch_mask[keep]
+                B = batch_ids.shape[0]
+                if B == 0:
+                    continue
                 for layer_idx in source_layers:
                     grad_mat = _average_jacobian_for_layer(
                         model,
@@ -328,6 +368,8 @@ def train_jacobian_lens(
                     J[layer_idx] += grad_mat
                 count += B
 
+            if count == 0:
+                raise JSpaceError("corpus has no valid (unmasked) positions")
             for layer in source_layers:
                 J[layer] = J[layer] / count
 
diff --git a/jspace/utils.py b/jspace/utils.py
index 9f57dc6..4d05b16 100644
--- a/jspace/utils.py
+++ b/jspace/utils.py
@@ -75,6 +75,26 @@ def hash_file(path: Path | str) -> str:
     return _compute_checksum(Path(path))
 
 
+def resolve_model_revision(model_name: str, revision: str | None) -> str | None:
+    """Resolve a model name to a pinned hub commit sha for cache keys.
+
+    Returns ``revision`` unchanged when provided. Local directory paths and
+    any hub failure (network offline, repo not found, huggingface_hub not
+    installed) resolve to None rather than raising.
+    """
+    if revision is not None:
+        return revision
+    if Path(model_name).is_dir():
+        return None
+    try:
+        from huggingface_hub import model_info
+
+        sha: str = model_info(model_name).sha
+        return sha
+    except Exception:
+        return None
+
+
 def _compute_checksum(path: Path) -> str:
     hasher = hashlib.sha256()
     with path.open("rb") as fh:
@@ -123,7 +143,7 @@ def load_lens_layer(cache_dir: Path, layer_idx: int) -> np.ndarray:
     matrix = data["J"]
     shape = tuple(int(dim) for dim in matrix.shape)
     _validate_shape(shape)
-    return cast(np.ndarray, np.asarray(matrix, dtype=np.float32))
+    return cast(np.ndarray, np.array(matrix))
 
 
 def lens_cache_exists(cache_dir: Path, layer_indices: list) -> bool:
diff --git a/scripts/train_lens.py b/scripts/train_lens.py
index f2ee1df..d84f23b 100644
--- a/scripts/train_lens.py
+++ b/scripts/train_lens.py
@@ -10,7 +10,7 @@
 from jspace import JSpaceError
 from jspace.jacobian_lens import train_jacobian_lens
 from jspace.model_adapter import layer_indices, load_model
-from jspace.utils import get_cache_dir, hash_file, model_fingerprint
+from jspace.utils import get_cache_dir, hash_file, model_fingerprint, resolve_model_revision
 from jspace.validation import validate_path, validate_workspace
 from jspace.viz import config_table, get_console, header_panel
 
@@ -103,6 +103,7 @@ def main():
             "need at least 2 to pick a penultimate target layer"
         )
     target_layer = args.target_layer if args.target_layer is not None else layers[-2]
+    resolved_rev = resolve_model_revision(args.model, args.model_revision)
     cache_dir = get_cache_dir(
         cache_base,
         model_fingerprint(
@@ -112,7 +113,7 @@ def main():
             max_positions=args.max_positions,
             dtype=args.dtype,
             output_dim_chunk=args.output_dim_chunk,
-            revision=args.model_revision,
+            revision=resolved_rev,
             corpus_hash=hash_file(corpus_path),
         ),
     )
diff --git a/scripts/workspace_geometry.py b/scripts/workspace_geometry.py
index 161ab12..235fe0e 100644
--- a/scripts/workspace_geometry.py
+++ b/scripts/workspace_geometry.py
@@ -21,7 +21,7 @@
 )
 from jspace.jacobian_lens import train_jacobian_lens
 from jspace.model_adapter import get_unembedding_matrix, layer_indices, load_model, normalize_fn
-from jspace.utils import get_cache_dir, hash_file, model_fingerprint
+from jspace.utils import get_cache_dir, hash_file, model_fingerprint, resolve_model_revision
 from jspace.validation import validate_path, validate_workspace
 from jspace.viz import (
     config_table,
@@ -48,18 +48,21 @@ def _fail(message: str) -> SystemExit:
 DEFAULT_PROBE_COUNT = 4096
 
 
-def _parse_probe_ids(probe_ids_str: str, vocab_size: int) -> list[int]:
+def _parse_probe_ids(probe_ids_str: str) -> list[int]:
     """Parse a comma-separated list of probe token ids."""
     try:
-        ids = [int(x.strip()) for x in probe_ids_str.split(",")]
+        return [int(x.strip()) for x in probe_ids_str.split(",")]
     except ValueError as exc:
         raise JSpaceError(
             f"invalid --probe-ids {probe_ids_str!r}: expected comma-separated integers"
         ) from exc
+
+
+def _check_probe_ids_in_range(ids: list[int], vocab_size: int) -> None:
+    """Reject probe token ids outside [0, vocab_size)."""
     for token_id in ids:
         if not 0 <= token_id < vocab_size:
             raise JSpaceError(f"probe token_id {token_id} out of range (vocab_size={vocab_size})")
-    return ids
 
 
 def _default_probe_ids(tokenizer, n: int) -> list[int]:
@@ -106,7 +109,7 @@ def compute_cka_block(V_by_layer: dict[int, torch.Tensor]) -> np.ndarray:
     """Compute layer-by-layer CKA similarity matrix."""
     layers = sorted(V_by_layer.keys())
     n = len(layers)
-    cka = np.zeros((n, n), dtype=np.float64)
+    cka: np.ndarray = np.zeros((n, n), dtype=np.float64)
     for i, li in enumerate(jl_track(layers, "CKA block")):
         for j, lj in enumerate(layers):
             if j < i:
@@ -152,6 +155,12 @@ def main() -> None:
     )
     args = parser.parse_args()
 
+    # Fail fast on malformed probe ids before any model loading or training.
+    try:
+        probe_ids_list = _parse_probe_ids(args.probe_ids) if args.probe_ids is not None else None
+    except JSpaceError as exc:
+        raise _fail(str(exc)) from exc
+
     dtype = getattr(torch, args.dtype)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -171,15 +180,33 @@ def main() -> None:
         revision=args.model_revision,
         allow_unlisted=args.allow_unlisted_model,
     )
-    corpus, attn_mask = _load_corpus(corpus_path, tokenizer, args.max_positions)
     layers = layer_indices(model)
     if args.target_layer is None and len(layers) < 2:
         raise _fail(
             f"{args.model} has {len(layers)} layer(s); "
             "need at least 2 to pick a penultimate target layer"
         )
+    if args.target_layer is not None and not 0 <= args.target_layer < len(layers):
+        raise _fail(
+            f"--target-layer {args.target_layer} out of range for {args.model} "
+            f"({len(layers)} layers); expected 0 <= --target-layer < {len(layers)}"
+        )
     target_layer = args.target_layer if args.target_layer is not None else layers[-2]
 
+    corpus, attn_mask = _load_corpus(corpus_path, tokenizer, args.max_positions)
+
+    W_U = get_unembedding_matrix(model)
+    vocab_size = W_U.shape[0]
+    if probe_ids_list is not None:
+        try:
+            _check_probe_ids_in_range(probe_ids_list, vocab_size)
+        except JSpaceError as exc:
+            raise _fail(str(exc)) from exc
+    else:
+        probe_ids_list = _default_probe_ids(tokenizer, args.n_probes)
+    probe_ids = torch.tensor(probe_ids_list, dtype=torch.long)
+
+    resolved_rev = resolve_model_revision(args.model, args.model_revision)
     cache_dir = get_cache_dir(
         cache_base,
         model_fingerprint(
@@ -189,7 +216,7 @@ def main() -> None:
             max_positions=args.max_positions,
             dtype=args.dtype,
             output_dim_chunk=args.output_dim_chunk,
-            revision=args.model_revision,
+            revision=resolved_rev,
             corpus_hash=hash_file(corpus_path),
         ),
     )
@@ -220,17 +247,6 @@ def main() -> None:
         attention_mask=attn_mask,
     )
 
-    W_U = get_unembedding_matrix(model)
-    vocab_size = W_U.shape[0]
-    try:
-        if args.probe_ids is not None:
-            probe_ids_list = _parse_probe_ids(args.probe_ids, vocab_size)
-        else:
-            probe_ids_list = _default_probe_ids(tokenizer, args.n_probes)
-    except JSpaceError as exc:
-        raise _fail(str(exc)) from exc
-    probe_ids = torch.tensor(probe_ids_list, dtype=torch.long)
-
     console.print(f"[bold]Building token geometry for[/] {len(probe_ids_list)} probe tokens")
     V_by_layer: dict[int, torch.Tensor] = {}
     for layer in jl_track(sorted(J.keys()), "Building V_l"):
@@ -254,6 +270,7 @@ def main() -> None:
         normalize_fn(model),
         layers=layers_with_v,
         probe_ids=probe_ids,
+        attention_mask=attn_mask,
     )
     start, end = infer_workspace_boundaries(discovery_metrics)
 
diff --git a/tests/test_bugfixes.py b/tests/test_bugfixes.py
index 1bc92a0..89323cb 100644
--- a/tests/test_bugfixes.py
+++ b/tests/test_bugfixes.py
@@ -1,5 +1,9 @@
 """Regression tests for bug fixes (see commit history)."""
 
+import json
+import subprocess
+import sys
+
 import numpy as np
 import pytest
 import torch
@@ -77,3 +81,165 @@ def test_apply_intervention_accepts_batched_inputs():
     )
     assert logits.shape[0] == 2
     assert logits.shape[1] == enc["input_ids"].shape[1]
+
+
+def test_inline_image_color_index_stays_in_palette():
+    """Near-white greys used to overflow to index 256 pre-fix (b05e14e)."""
+    from scripts.inline_image import _rgb_to_256
+
+    assert _rgb_to_256(255, 255, 255) == 231
+    assert _rgb_to_256(250, 250, 250) == 231
+    assert _rgb_to_256(248, 248, 248) == 231  # was 256 before the fix
+    assert _rgb_to_256(247, 247, 247) == 255
+    assert _rgb_to_256(255, 0, 0) == 196
+    for r in range(0, 256, 11):
+        for g in range(0, 256, 11):
+            for b in range(0, 256, 11):
+                assert 0 <= _rgb_to_256(r, g, b) <= 255
+
+
+def test_inline_image_render_escape_codes_bounded(tmp_path):
+    from PIL import Image
+
+    from scripts.inline_image import render
+
+    Image.new("RGB", (3, 2), (250, 250, 250)).save(tmp_path / "near_white.png")
+    out = render(str(tmp_path / "near_white.png"), width=3)
+    for seq in out.split("\033[38;5;")[1:]:
+        assert 0 <= int(seq.split("m")[0]) <= 255
+
+
+def test_inline_image_closes_image_handle(tmp_path, monkeypatch):
+    from PIL import Image
+
+    from scripts import inline_image
+
+    Image.new("RGB", (2, 2), (10, 20, 30)).save(tmp_path / "t.png")
+    opened = []
+    real_open = Image.open
+
+    def spy(path):
+        handle = real_open(path)
+        opened.append(handle)
+        return handle
+
+    monkeypatch.setattr(inline_image.Image, "open", spy)
+    inline_image.render(str(tmp_path / "t.png"), width=2)
+
+    assert len(opened) == 1
+    assert opened[0].fp is None  # closed by the ``with`` block
+
+
+def _wg_argv(tmp_path, corpus, *extra):
+    return [
+        "prog",
+        "--model",
+        "sshleifer/tiny-gpt2",
+        "--corpus",
+        str(corpus),
+        "--workspace",
+        str(tmp_path),
+        "--cache-dir",
+        str(tmp_path / "cache"),
+        "--output-dir",
+        str(tmp_path / "out"),
+        *extra,
+    ]
+
+
+def _tiny_gpt2(n_layer: int = 2) -> torch.nn.Module:
+    """In-process GPT-2 stand-in so CLI tests never hit the hub."""
+    from transformers import GPT2Config, GPT2LMHeadModel
+
+    config = GPT2Config(
+        n_layer=n_layer, n_head=1, n_embd=8, n_positions=8, n_ctx=8, vocab_size=16
+    )
+    return GPT2LMHeadModel(config)
+
+
+def test_workspace_geometry_probe_ids_error_precedes_model_load(tmp_path):
+    """Malformed --probe-ids must fail before model loading or training."""
+    corpus = tmp_path / "corpus.json"
+    corpus.write_text(json.dumps(["The cat sat."]))
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-m",
+            "scripts.workspace_geometry",
+            "--model",
+            "totally-bogus-model-name",
+            "--corpus",
+            str(corpus),
+            "--workspace",
+            str(tmp_path),
+            "--cache-dir",
+            str(tmp_path / "cache"),
+            "--output-dir",
+            str(tmp_path / "out"),
+            "--probe-ids",
+            "1,2,x",
+        ],
+        capture_output=True,
+        text=True,
+    )
+    assert result.returncode == 1
+    assert "--probe-ids" in result.stderr
+    # The allowlist error would name the model; it must never be reached.
+    assert "totally-bogus-model-name" not in result.stderr
+
+
+def test_workspace_geometry_rejects_single_layer_model(tmp_path, monkeypatch, capsys):
+    import scripts.workspace_geometry as wg
+
+    monkeypatch.setattr(wg, "load_model", lambda *a, **k: (_tiny_gpt2(n_layer=1), object()))
+    corpus = tmp_path / "corpus.json"
+    corpus.write_text(json.dumps(["The cat sat."]))
+    monkeypatch.setattr(sys, "argv", _wg_argv(tmp_path, corpus))
+
+    with pytest.raises(SystemExit) as exc:
+        wg.main()
+    assert exc.value.code == 1
+    assert "1 layer" in capsys.readouterr().err
+
+
+@pytest.mark.parametrize("bad_layer", ["-1", "99"])
+def test_workspace_geometry_rejects_out_of_range_target_layer(
+    tmp_path, monkeypatch, capsys, bad_layer
+):
+    import scripts.workspace_geometry as wg
+
+    def boom(*a, **k):
+        raise AssertionError("training must not run for an invalid --target-layer")
+
+    monkeypatch.setattr(wg, "train_jacobian_lens", boom)
+    monkeypatch.setattr(wg, "load_model", lambda *a, **k: (_tiny_gpt2(), object()))
+    corpus = tmp_path / "corpus.json"
+    corpus.write_text(json.dumps(["The cat sat on the mat."]))
+    monkeypatch.setattr(sys, "argv", _wg_argv(tmp_path, corpus, "--target-layer", bad_layer))
+
+    with pytest.raises(SystemExit) as exc:
+        wg.main()
+    assert exc.value.code == 1
+    assert "out of range" in capsys.readouterr().err
+
+
+def test_workspace_geometry_rejects_out_of_range_probe_ids_before_training(
+    tmp_path, monkeypatch, capsys
+):
+    import scripts.workspace_geometry as wg
+
+    def boom(*a, **k):
+        raise AssertionError("training must not run for out-of-range --probe-ids")
+
+    monkeypatch.setattr(wg, "train_jacobian_lens", boom)
+    monkeypatch.setattr(wg, "load_model", lambda *a, **k: (_tiny_gpt2(), object()))
+    ids = torch.ones(1, 4, dtype=torch.long)
+    monkeypatch.setattr(wg, "_load_corpus", lambda *a, **k: (ids, ids.clone()))
+    corpus = tmp_path / "corpus.json"
+    corpus.write_text(json.dumps(["The cat sat on the mat."]))
+    monkeypatch.setattr(sys, "argv", _wg_argv(tmp_path, corpus, "--probe-ids", f"1,{10**9}"))
+
+    with pytest.raises(SystemExit) as exc:
+        wg.main()
+    assert exc.value.code == 1
+    assert "out of range" in capsys.readouterr().err
diff --git a/tests/test_decomposition.py b/tests/test_decomposition.py
index 5240f28..7b6b6eb 100644
--- a/tests/test_decomposition.py
+++ b/tests/test_decomposition.py
@@ -20,11 +20,15 @@ def test_decomposition_accepts_j_l_and_w_u():
     d = 64
     vocab = 100
     h = torch.randn(d)
-    J_l = np.eye(d).astype(np.float32)
+    # Non-symmetric J_l: only the W_U @ J_l (rows = J-lens vectors) geometry
+    # matches the reference decomposition below; W_U @ J_l.T would not.
+    J_l = (np.arange(d * d, dtype=np.float64).reshape(d, d) / d**2).astype(np.float32)
     W_U = torch.randn(vocab, d)
     a, h_J, h_perp = decompose_jspace(h, J_l=J_l, W_U=W_U, k=10)
     assert a.shape == (vocab,)
     assert torch.allclose(h, h_J + h_perp, atol=1e-4)
+    _, h_J_ref, _ = decompose_jspace(h, V=W_U @ torch.from_numpy(J_l), k=10)
+    assert torch.allclose(h_J, h_J_ref, atol=1e-4)
 
 
 def test_non_negative_decomposition_returns_non_negative_coefficients():
@@ -88,3 +92,38 @@ def test_non_negative_coeffs_match_nnls_refit():
         torch.from_numpy(expected).float(),
         atol=1e-4,
     )
+
+
+def test_occupancy_returns_zero_for_zero_hidden_state():
+    """h = 0 has an empty support; occupancy must stop at k=0, not max_k."""
+    d = 32
+    V = torch.randn(50, d)
+    h = torch.zeros(d)
+    k = jspace_occupancy(h, V=V, max_k=10, threshold=0.0, random_seed=0)
+    assert k == 0
+
+
+def test_occupancy_stops_when_signal_is_exactly_represented():
+    """A 2-sparse h must stop at k=2 even at threshold 0 (residual is 0)."""
+    d = 32
+    V = torch.zeros(50, d)
+    V[3, 0] = 2.0
+    V[17, 1] = 1.5
+    h = V[3] + V[17]
+    k = jspace_occupancy(h, V=V, max_k=10, threshold=0.0, random_seed=0)
+    assert k == 2
+
+
+def test_build_V_l_uses_j_lens_vector_geometry():
+    """scripts.build_V_l rows must be W_U @ J_l, not W_U @ J_l.T."""
+    from scripts.workspace_geometry import build_V_l
+
+    d = 8
+    vocab = 5
+    J_l = np.arange(d * d, dtype=np.float32).reshape(d, d) / d**2  # non-symmetric
+    W_U = torch.randn(vocab, d)
+    probe_ids = torch.tensor([0, 2, 4])
+    V = build_V_l(J_l, W_U, probe_ids)
+    expected = W_U[probe_ids] @ torch.from_numpy(J_l)
+    expected = expected - expected.mean(dim=0, keepdim=True)
+    assert torch.allclose(V, expected)
diff --git a/tests/test_discovery.py b/tests/test_discovery.py
index 2714397..f199707 100644
--- a/tests/test_discovery.py
+++ b/tests/test_discovery.py
@@ -1,7 +1,15 @@
+"""Tests for workspace discovery: CKA/boundaries and attention-mask validation."""
+
 import numpy as np
+import pytest
 import torch
 
-from jspace.discovery import centered_kernel_alignment, infer_workspace_boundaries
+from jspace import JSpaceError
+from jspace.discovery import (
+    centered_kernel_alignment,
+    compute_discovery_metrics,
+    infer_workspace_boundaries,
+)
 
 
 def test_cka_and_boundary_inference():
@@ -17,3 +25,20 @@ def test_cka_and_boundary_inference():
     }
     start, end = infer_workspace_boundaries(metrics)
     assert 0 <= start <= end < 10
+
+
+def test_discovery_metrics_rejects_mismatched_mask():
+    model = torch.nn.Linear(4, 4)
+    corpus = torch.ones(1, 3, dtype=torch.long)
+    bad_mask = torch.ones(1, 4, dtype=torch.long)
+    with pytest.raises(JSpaceError, match="attention_mask"):
+        compute_discovery_metrics(
+            model,
+            object(),
+            {},
+            corpus,
+            torch.ones(4, 4),
+            lambda x: x,
+            layers=[0],
+            attention_mask=bad_mask,
+        )
diff --git a/tests/test_jacobian_lens.py b/tests/test_jacobian_lens.py
index c06fc2d..212adb2 100644
--- a/tests/test_jacobian_lens.py
+++ b/tests/test_jacobian_lens.py
@@ -1,6 +1,8 @@
+import pytest
 import torch
 
 from jspace.jacobian_lens import (
+    _attach_frozen_qk_hooks,
     _average_jacobian_for_layer,
     _base_model,
     _capture_h_l,
@@ -126,3 +128,8 @@ def mock_run_from_layer(model, layer_idx, h_l, input_ids, attention_mask, target
 
     expected = torch.eye(d_model, dtype=torch.float64) * expected_diag
     assert torch.allclose(J, expected, atol=1e-4)
+
+
+def test_frozen_qk_warns_when_no_projection_matches():
+    with pytest.warns(UserWarning, match="no query/key projections"):
+        _attach_frozen_qk_hooks(torch.nn.Linear(4, 4))
diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 0000000..7e378c4
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,60 @@
+"""Tests for jspace.utils cache IO, fingerprints, and revision resolution."""
+
+import numpy as np
+import pytest
+
+from jspace.utils import (
+    load_lens_layer,
+    resolve_model_revision,
+    save_lens_layer,
+)
+
+
+def test_load_lens_layer_preserves_saved_dtype(tmp_path):
+    matrix = np.random.rand(8, 8).astype(np.float64)
+    save_lens_layer(tmp_path, 0, matrix)
+    loaded = load_lens_layer(tmp_path, 0)
+    assert isinstance(loaded, np.ndarray)
+    assert loaded.dtype == np.float64
+    assert np.array_equal(loaded, matrix)
+
+
+def test_load_lens_layer_returns_defensive_copy(tmp_path):
+    matrix = np.random.rand(8, 8).astype(np.float32)
+    save_lens_layer(tmp_path, 0, matrix)
+    loaded = load_lens_layer(tmp_path, 0)
+    loaded[0, 0] = -123.0
+    reloaded = load_lens_layer(tmp_path, 0)
+    assert reloaded[0, 0] == pytest.approx(matrix[0, 0])
+
+
+def test_resolve_model_revision_passes_through_explicit_revision():
+    assert resolve_model_revision("gpt2", "abc123") == "abc123"
+
+
+def test_resolve_model_revision_returns_none_for_local_dir(tmp_path, monkeypatch):
+    def _boom(repo_id):  # pragma: no cover - must never be called
+        raise AssertionError("model_info must not be hit for local paths")
+
+    monkeypatch.setattr("huggingface_hub.model_info", _boom)
+    assert resolve_model_revision(str(tmp_path), None) is None
+
+
+def test_resolve_model_revision_returns_none_on_hub_failure(monkeypatch):
+    def _raise(repo_id):
+        raise RuntimeError("network is unreachable")
+
+    monkeypatch.setattr("huggingface_hub.model_info", _raise)
+    assert resolve_model_revision("bogus-repo-that-cannot-resolve", None) is None
+
+
+def test_resolve_model_revision_returns_hub_sha(monkeypatch):
+    class _Info:
+        sha = "cafe0123beef"
+
+    def _fake_model_info(repo_id):
+        assert repo_id == "gpt2"
+        return _Info()
+
+    monkeypatch.setattr("huggingface_hub.model_info", _fake_model_info)
+    assert resolve_model_revision("gpt2", None) == "cafe0123beef"