fix: bug-fix review follow-ups and regression test backfill

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"],
 )
 ```
 

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"

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.

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,16 +42,29 @@ 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.
 
     If `layers` is omitted, all model layers are used. Pass a subset when the
     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,28 +89,48 @@ 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()
         acc[idx] = float(correct / (valid.sum() + 1e-9))
 
         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}
 

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,14 +94,47 @@ 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,
     attention_mask: torch.Tensor,
     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
 

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:

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),
         ),
     )