From 15a2528a38d8f2c24716bbed982f4846609b0042 Mon Sep 17 00:00:00 2001 From: J-Space Developer Date: Wed, 16 Sep 2026 06:12:42 -0500 Subject: [PATCH 1/2] fix: address deep-review findings on b05e14e and backfill regression tests Findings from a full review of the bug-fix commits, now fixed: - discovery.compute_discovery_metrics accepts and forwards attention_mask; kurtosis/autocorr are computed over unmasked positions only (the same pad-corruption class fixed in train_jacobian_lens had survived here) - jspace_occupancy stops when the residual is exactly exhausted instead of running to max_k (h=0 previously returned max_k) - train_jacobian_lens drops fully-masked corpus rows from the J_l average and raises when no valid positions remain, instead of caching zeros - position ids derive from the attention mask, fixing the left-padding IndexError (get_position_ids yielded -1 for leading pads) - _capture_h_l captures the pre-final-norm residual for the last source layer via a norm-input hook (transformers >= 5.x ties hidden_states[-1] to the post-norm last_hidden_state) - load_lens_layer restores dtype preservation and the defensive copy (0ceefe1 had silently forced float32 via asarray) - add utils.resolve_model_revision so cache fingerprints key on the resolved hub sha rather than the often-None --model-revision - workspace-geometry CLI validates --probe-ids before model load and validates --target-layer range after load - README example rewritten against the current API; demo.ipynb forwards the attention mask to train_jacobian_lens and compute_discovery_metrics Backfilled regression tests (all verified to fail on pre-fix code where applicable): pad-exclusion end-to-end, left-pad equivalence, pre-norm last-layer capture, single-forward and requires_grad restoration, frozen_qk gradient/warning behavior, occupancy zero-residual stop, build_V_l geometry, load_lens_layer dtype/copy semantics, inline_image palette overflow and handle close, fast-fail CLI validation, discovery masking. Full suite 80 passed; mypy and ruff clean. --- README.md | 73 ++++++++------- demo.ipynb | 6 +- jspace/decomposition.py | 3 + jspace/discovery.py | 47 ++++++++-- jspace/jacobian_lens.py | 50 +++++++++- jspace/utils.py | 22 ++++- scripts/train_lens.py | 5 +- scripts/workspace_geometry.py | 53 +++++++---- tests/test_bugfixes.py | 155 +++++++++++++++++++++++++++++++ tests/test_decomposition.py | 41 ++++++++- tests/test_discovery.py | 106 +++++++++++++++++++++- tests/test_jacobian_lens.py | 166 ++++++++++++++++++++++++++++++++++ tests/test_utils.py | 60 ++++++++++++ 13 files changed, 717 insertions(+), 70 deletions(-) create mode 100644 tests/test_utils.py 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..bcbcc2e 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,154 @@ 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 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): + from transformers import GPT2Config, GPT2LMHeadModel + + import scripts.workspace_geometry as wg + + config = GPT2Config(n_layer=1, n_head=1, n_embd=8, n_positions=8, n_ctx=8, vocab_size=16) + monkeypatch.setattr(wg, "load_model", lambda *a, **k: (GPT2LMHeadModel(config), 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) + 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) + 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..88ade07 100644 --- a/tests/test_discovery.py +++ b/tests/test_discovery.py @@ -1,7 +1,21 @@ +"""Tests for workspace discovery: CKA/boundaries and attention-mask handling.""" + 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, +) +from jspace.model_adapter import ( + get_unembedding_matrix, + layer_indices, + load_model, + normalize_fn, +) def test_cka_and_boundary_inference(): @@ -17,3 +31,93 @@ def test_cka_and_boundary_inference(): } start, end = infer_workspace_boundaries(metrics) assert 0 <= start <= end < 10 + + +def _tiny_jacobians(model) -> dict[int, np.ndarray]: + """Small random J_l matrices, one per layer (metric values are irrelevant).""" + d = model.config.n_embd + rng = np.random.default_rng(0) + return { + layer: (0.01 * np.eye(d) + rng.normal(scale=0.01, size=(d, d))).astype(np.float32) + for layer in layer_indices(model) + } + + +@pytest.fixture() +def tiny_gpt2(): + return load_model("sshleifer/tiny-gpt2", torch.device("cpu"), torch.float32) + + +def _pad(ids: torch.Tensor, tokenizer, n_pads: int) -> tuple[torch.Tensor, torch.Tensor]: + """Append ``n_pads`` trailing pad columns plus the matching attention mask.""" + pad_id = tokenizer.pad_token_id + if pad_id is None: + pad_id = tokenizer.eos_token_id + padded = torch.cat( + [ids, torch.full((ids.shape[0], n_pads), pad_id, dtype=ids.dtype)], dim=1 + ) + mask = torch.cat( + [torch.ones_like(ids), torch.zeros(ids.shape[0], n_pads, dtype=ids.dtype)], dim=1 + ) + return padded, mask + + +def _metrics(model, tokenizer, corpus, attention_mask=None): + return compute_discovery_metrics( + model, + tokenizer, + _tiny_jacobians(model), + corpus, + get_unembedding_matrix(model), + normalize_fn(model), + attention_mask=attention_mask, + ) + + +def test_masked_metrics_ignore_trailing_pads(tiny_gpt2): + """Metrics with a mask must be identical for a row with any number of pads.""" + model, tokenizer = tiny_gpt2 + ids = tokenizer("The cat sat on the mat.", return_tensors="pt")["input_ids"] + + reference = _metrics(model, tokenizer, ids) + padded2, mask2 = _pad(ids, tokenizer, 2) + padded5, mask5 = _pad(ids, tokenizer, 5) + m2 = _metrics(model, tokenizer, padded2, attention_mask=mask2) + m5 = _metrics(model, tokenizer, padded5, attention_mask=mask5) + + for key in ("kurtosis", "accuracy", "autocorr"): + assert np.allclose(reference[key], m2[key], atol=1e-5), key + assert np.allclose(reference[key], m5[key], atol=1e-5), key + + # Without the mask the pad positions pollute at least one metric. + unmasked = _metrics(model, tokenizer, padded2) + assert any( + not np.allclose(reference[key], unmasked[key], atol=1e-5) + for key in ("kurtosis", "accuracy", "autocorr") + ) + + +def test_discovery_metrics_forwards_attention_mask(tiny_gpt2, monkeypatch): + model, tokenizer = tiny_gpt2 + ids = tokenizer("The cat sat.", return_tensors="pt")["input_ids"] + mask = torch.ones_like(ids) + captured = {} + orig_forward = model.forward + + def spy(*args, **kwargs): + captured.update(kwargs) + return orig_forward(*args, **kwargs) + + monkeypatch.setattr(model, "forward", spy) + _metrics(model, tokenizer, ids, attention_mask=mask) + + assert "attention_mask" in captured + assert torch.equal(captured["attention_mask"], mask) + + +def test_discovery_metrics_rejects_mismatched_mask(tiny_gpt2): + model, tokenizer = tiny_gpt2 + ids = tokenizer("The cat sat.", return_tensors="pt")["input_ids"] + bad_mask = torch.ones(1, ids.shape[1] + 1, dtype=torch.long) + with pytest.raises(JSpaceError, match="attention_mask"): + _metrics(model, tokenizer, ids, attention_mask=bad_mask) diff --git a/tests/test_jacobian_lens.py b/tests/test_jacobian_lens.py index c06fc2d..0f8ce4d 100644 --- a/tests/test_jacobian_lens.py +++ b/tests/test_jacobian_lens.py @@ -1,6 +1,9 @@ +import numpy as np +import pytest import torch from jspace.jacobian_lens import ( + _attach_frozen_qk_hooks, _average_jacobian_for_layer, _base_model, _capture_h_l, @@ -126,3 +129,166 @@ 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_fully_masked_rows_are_excluded_from_j_lens(): + """A fully-masked row must not dilute J_l: result equals the corpus without it.""" + model, tokenizer = load_model("sshleifer/tiny-gpt2", torch.device("cpu"), torch.float32) + prompts = ["The cat sat.", "In 1950, scientists discovered."] + enc = tokenizer(prompts, return_tensors="pt", padding=True, truncation=True, max_length=16) + corpus = enc["input_ids"] + mask = enc["attention_mask"] + full_pad = torch.full_like(corpus[0], tokenizer.pad_token_id) + corpus_with_pad = torch.cat([corpus, full_pad.unsqueeze(0)], dim=0) + mask_with_pad = torch.cat([mask, torch.zeros_like(mask[0]).unsqueeze(0)], dim=0) + + target_layer = layer_indices(model)[-1] + common = { + "target_layer": target_layer, + "dtype": torch.float32, + "max_positions": 16, + "batch_size": 2, + "output_dim_chunk": 16, + } + J_ref = train_jacobian_lens(model, corpus, attention_mask=mask, **common) + J_test = train_jacobian_lens( + model, corpus_with_pad, attention_mask=mask_with_pad, **common + ) + for layer in J_ref: + assert np.allclose(J_ref[layer], J_test[layer], atol=1e-6) + + +def test_left_padded_batch_matches_right_padded_jacobian(): + """Left padding must not crash, and real tokens must give the same J_l.""" + model, tokenizer = load_model("sshleifer/tiny-gpt2", torch.device("cpu"), torch.float32) + ids = tokenizer(["The cat sat."], return_tensors="pt")["input_ids"][0] + T = ids.numel() + n_pad = 2 + pad_id = tokenizer.pad_token_id + target_layer = layer_indices(model)[-1] + + def jacobian(input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: + return _average_jacobian_for_layer( + model, + input_ids.unsqueeze(0), + attention_mask.unsqueeze(0), + layer_idx=0, + output_dim_chunk=16, + target_layer=target_layer, + ) + + right_ids = torch.cat([ids, torch.full((n_pad,), pad_id, dtype=torch.long)]) + right_mask = torch.cat([torch.ones(T, dtype=torch.long), torch.zeros(n_pad, dtype=torch.long)]) + left_ids = torch.cat([torch.full((n_pad,), pad_id, dtype=torch.long), ids]) + left_mask = torch.cat([torch.zeros(n_pad, dtype=torch.long), torch.ones(T, dtype=torch.long)]) + + J_right = jacobian(right_ids, right_mask) + J_left = jacobian(left_ids, left_mask) + assert torch.allclose(J_left, J_right, atol=1e-4) + + +def test_capture_h_l_last_layer_is_pre_norm_residual(): + """transformers >= 5.x ties hidden_states[-1] to the post-norm residual; h_l must not be.""" + model, tokenizer = load_model("sshleifer/tiny-gpt2", torch.device("cpu"), torch.float32) + input_ids = tokenizer(["The cat sat."], return_tensors="pt")["input_ids"] + attention_mask = torch.ones_like(input_ids) + last = layer_indices(model)[-1] + base = _base_model(model) + + block_out: list[torch.Tensor] = [] + + def capture_block_output(module, input_tuple, output): + block_out.append((output[0] if isinstance(output, tuple) else output).detach().clone()) + + handle = base.h[last].register_forward_hook(capture_block_output) + try: + base( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=get_position_ids(attention_mask), + return_dict=True, + ) + finally: + handle.remove() + + h_l = _capture_h_l(model, input_ids, attention_mask, last) + assert torch.allclose(h_l, block_out[0], atol=1e-6) + assert not torch.allclose(h_l, base.ln_f(block_out[0]), atol=1e-4) + + # Train-path preservation: with target_layer=None the last layer maps to itself. + z = _run_from_layer(model, last, h_l, input_ids, attention_mask, target_layer=None) + assert torch.equal(z, h_l) + + +def test_target_layer_none_runs_single_forward(): + """target_layer=None must not re-run the model (one embedding forward per pass).""" + model, tokenizer = load_model("sshleifer/tiny-gpt2", torch.device("cpu"), torch.float32) + input_ids = tokenizer(["The cat sat."], return_tensors="pt")["input_ids"] + attention_mask = torch.ones_like(input_ids) + wte = _base_model(model).wte + calls = {"n": 0} + handle = wte.register_forward_hook( + lambda module, inp, out: calls.__setitem__("n", calls["n"] + 1) + ) + try: + h_l = _capture_h_l(model, input_ids, attention_mask, layer_idx=0) + calls["n"] = 0 + _run_from_layer(model, 0, h_l, input_ids, attention_mask, target_layer=None) + assert calls["n"] == 1 + + calls["n"] = 0 + train_jacobian_lens( + model, + input_ids, + target_layer=None, + dtype=torch.float32, + batch_size=1, + output_dim_chunk=16, + ) + # capture(l0) + run(l0) + capture(l1); l1 short-circuits before any forward. + assert calls["n"] == 3 + finally: + handle.remove() + + +def test_vjp_restores_requires_grad_flags(): + model, tokenizer = load_model("sshleifer/tiny-gpt2", torch.device("cpu"), torch.float32) + input_ids = tokenizer(["The cat sat."], return_tensors="pt")["input_ids"] + attention_mask = torch.ones_like(input_ids) + params = list(model.parameters()) + for i, p in enumerate(params): + p.requires_grad_(i % 2 == 0) + saved = [p.requires_grad for p in params] + try: + _average_jacobian_for_layer( + model, input_ids, attention_mask, layer_idx=0, output_dim_chunk=16 + ) + assert [p.requires_grad for p in params] == saved + finally: + for p in params: + p.requires_grad_(True) + + +def test_frozen_qk_detaches_gpt2_c_attn_qk(): + """Gradients must not flow into c_attn's Q/K weight columns; V columns must still get grad.""" + model, _ = load_model("sshleifer/tiny-gpt2", torch.device("cpu"), torch.float32) + n_embd = model.config.n_embd + c_attn = model.transformer.h[0].attn.c_attn + hidden = torch.randn(1, 4, n_embd, requires_grad=True) + handles = _attach_frozen_qk_hooks(model) + try: + out = c_attn(hidden) + out.backward(torch.ones_like(out)) + finally: + for h in handles: + h.remove() + grad = c_attn.weight.grad + assert grad is not None + # Conv1D weight is (nx, nf); columns [0:n]=Q, [n:2n]=K, [2n:3n]=V. + assert grad[:, : 2 * n_embd].abs().max().item() == 0.0 + assert grad[:, 2 * n_embd :].abs().max().item() > 0.0 + + +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" From 063add148af7ab8130bcde7d84b62336cc04b489 Mon Sep 17 00:00:00 2001 From: J-Space Developer Date: Wed, 16 Sep 2026 07:38:50 -0500 Subject: [PATCH 2/2] test: drop tiny-gpt2-loading regression tests to keep the suite fast Removes the backfilled tests that load sshleifer/tiny-gpt2 (suite was ~4 min; now 72 passed in ~64s, back under the pre-backfill baseline): - test_jacobian_lens: pad-exclusion, left-pad equivalence, pre-norm capture, single-forward, requires_grad restore, frozen_qk gradient (kept the model-free frozen_qk warning test) - test_discovery: mask-invariance and mask-forwarding tests plus the tiny_gpt2 fixture (kept the mask shape-validation test, now running against a dummy module) - test_bugfixes: the out-of-range --target-layer / --probe-ids CLI tests now stub load_model/_load_corpus with an in-process GPT-2 stand-in instead of hitting the hub, keeping the validation regression coverage --- tests/test_bugfixes.py | 19 ++++- tests/test_discovery.py | 109 ++++-------------------- tests/test_jacobian_lens.py | 159 ------------------------------------ 3 files changed, 30 insertions(+), 257 deletions(-) diff --git a/tests/test_bugfixes.py b/tests/test_bugfixes.py index bcbcc2e..89323cb 100644 --- a/tests/test_bugfixes.py +++ b/tests/test_bugfixes.py @@ -147,6 +147,16 @@ def _wg_argv(tmp_path, corpus, *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" @@ -179,12 +189,9 @@ def test_workspace_geometry_probe_ids_error_precedes_model_load(tmp_path): def test_workspace_geometry_rejects_single_layer_model(tmp_path, monkeypatch, capsys): - from transformers import GPT2Config, GPT2LMHeadModel - import scripts.workspace_geometry as wg - config = GPT2Config(n_layer=1, n_head=1, n_embd=8, n_positions=8, n_ctx=8, vocab_size=16) - monkeypatch.setattr(wg, "load_model", lambda *a, **k: (GPT2LMHeadModel(config), object())) + 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)) @@ -205,6 +212,7 @@ 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)) @@ -224,6 +232,9 @@ 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}")) diff --git a/tests/test_discovery.py b/tests/test_discovery.py index 88ade07..f199707 100644 --- a/tests/test_discovery.py +++ b/tests/test_discovery.py @@ -1,4 +1,4 @@ -"""Tests for workspace discovery: CKA/boundaries and attention-mask handling.""" +"""Tests for workspace discovery: CKA/boundaries and attention-mask validation.""" import numpy as np import pytest @@ -10,12 +10,6 @@ compute_discovery_metrics, infer_workspace_boundaries, ) -from jspace.model_adapter import ( - get_unembedding_matrix, - layer_indices, - load_model, - normalize_fn, -) def test_cka_and_boundary_inference(): @@ -33,91 +27,18 @@ def test_cka_and_boundary_inference(): assert 0 <= start <= end < 10 -def _tiny_jacobians(model) -> dict[int, np.ndarray]: - """Small random J_l matrices, one per layer (metric values are irrelevant).""" - d = model.config.n_embd - rng = np.random.default_rng(0) - return { - layer: (0.01 * np.eye(d) + rng.normal(scale=0.01, size=(d, d))).astype(np.float32) - for layer in layer_indices(model) - } - - -@pytest.fixture() -def tiny_gpt2(): - return load_model("sshleifer/tiny-gpt2", torch.device("cpu"), torch.float32) - - -def _pad(ids: torch.Tensor, tokenizer, n_pads: int) -> tuple[torch.Tensor, torch.Tensor]: - """Append ``n_pads`` trailing pad columns plus the matching attention mask.""" - pad_id = tokenizer.pad_token_id - if pad_id is None: - pad_id = tokenizer.eos_token_id - padded = torch.cat( - [ids, torch.full((ids.shape[0], n_pads), pad_id, dtype=ids.dtype)], dim=1 - ) - mask = torch.cat( - [torch.ones_like(ids), torch.zeros(ids.shape[0], n_pads, dtype=ids.dtype)], dim=1 - ) - return padded, mask - - -def _metrics(model, tokenizer, corpus, attention_mask=None): - return compute_discovery_metrics( - model, - tokenizer, - _tiny_jacobians(model), - corpus, - get_unembedding_matrix(model), - normalize_fn(model), - attention_mask=attention_mask, - ) - - -def test_masked_metrics_ignore_trailing_pads(tiny_gpt2): - """Metrics with a mask must be identical for a row with any number of pads.""" - model, tokenizer = tiny_gpt2 - ids = tokenizer("The cat sat on the mat.", return_tensors="pt")["input_ids"] - - reference = _metrics(model, tokenizer, ids) - padded2, mask2 = _pad(ids, tokenizer, 2) - padded5, mask5 = _pad(ids, tokenizer, 5) - m2 = _metrics(model, tokenizer, padded2, attention_mask=mask2) - m5 = _metrics(model, tokenizer, padded5, attention_mask=mask5) - - for key in ("kurtosis", "accuracy", "autocorr"): - assert np.allclose(reference[key], m2[key], atol=1e-5), key - assert np.allclose(reference[key], m5[key], atol=1e-5), key - - # Without the mask the pad positions pollute at least one metric. - unmasked = _metrics(model, tokenizer, padded2) - assert any( - not np.allclose(reference[key], unmasked[key], atol=1e-5) - for key in ("kurtosis", "accuracy", "autocorr") - ) - - -def test_discovery_metrics_forwards_attention_mask(tiny_gpt2, monkeypatch): - model, tokenizer = tiny_gpt2 - ids = tokenizer("The cat sat.", return_tensors="pt")["input_ids"] - mask = torch.ones_like(ids) - captured = {} - orig_forward = model.forward - - def spy(*args, **kwargs): - captured.update(kwargs) - return orig_forward(*args, **kwargs) - - monkeypatch.setattr(model, "forward", spy) - _metrics(model, tokenizer, ids, attention_mask=mask) - - assert "attention_mask" in captured - assert torch.equal(captured["attention_mask"], mask) - - -def test_discovery_metrics_rejects_mismatched_mask(tiny_gpt2): - model, tokenizer = tiny_gpt2 - ids = tokenizer("The cat sat.", return_tensors="pt")["input_ids"] - bad_mask = torch.ones(1, ids.shape[1] + 1, dtype=torch.long) +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"): - _metrics(model, tokenizer, ids, attention_mask=bad_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 0f8ce4d..212adb2 100644 --- a/tests/test_jacobian_lens.py +++ b/tests/test_jacobian_lens.py @@ -1,4 +1,3 @@ -import numpy as np import pytest import torch @@ -131,164 +130,6 @@ def mock_run_from_layer(model, layer_idx, h_l, input_ids, attention_mask, target assert torch.allclose(J, expected, atol=1e-4) -def test_fully_masked_rows_are_excluded_from_j_lens(): - """A fully-masked row must not dilute J_l: result equals the corpus without it.""" - model, tokenizer = load_model("sshleifer/tiny-gpt2", torch.device("cpu"), torch.float32) - prompts = ["The cat sat.", "In 1950, scientists discovered."] - enc = tokenizer(prompts, return_tensors="pt", padding=True, truncation=True, max_length=16) - corpus = enc["input_ids"] - mask = enc["attention_mask"] - full_pad = torch.full_like(corpus[0], tokenizer.pad_token_id) - corpus_with_pad = torch.cat([corpus, full_pad.unsqueeze(0)], dim=0) - mask_with_pad = torch.cat([mask, torch.zeros_like(mask[0]).unsqueeze(0)], dim=0) - - target_layer = layer_indices(model)[-1] - common = { - "target_layer": target_layer, - "dtype": torch.float32, - "max_positions": 16, - "batch_size": 2, - "output_dim_chunk": 16, - } - J_ref = train_jacobian_lens(model, corpus, attention_mask=mask, **common) - J_test = train_jacobian_lens( - model, corpus_with_pad, attention_mask=mask_with_pad, **common - ) - for layer in J_ref: - assert np.allclose(J_ref[layer], J_test[layer], atol=1e-6) - - -def test_left_padded_batch_matches_right_padded_jacobian(): - """Left padding must not crash, and real tokens must give the same J_l.""" - model, tokenizer = load_model("sshleifer/tiny-gpt2", torch.device("cpu"), torch.float32) - ids = tokenizer(["The cat sat."], return_tensors="pt")["input_ids"][0] - T = ids.numel() - n_pad = 2 - pad_id = tokenizer.pad_token_id - target_layer = layer_indices(model)[-1] - - def jacobian(input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - return _average_jacobian_for_layer( - model, - input_ids.unsqueeze(0), - attention_mask.unsqueeze(0), - layer_idx=0, - output_dim_chunk=16, - target_layer=target_layer, - ) - - right_ids = torch.cat([ids, torch.full((n_pad,), pad_id, dtype=torch.long)]) - right_mask = torch.cat([torch.ones(T, dtype=torch.long), torch.zeros(n_pad, dtype=torch.long)]) - left_ids = torch.cat([torch.full((n_pad,), pad_id, dtype=torch.long), ids]) - left_mask = torch.cat([torch.zeros(n_pad, dtype=torch.long), torch.ones(T, dtype=torch.long)]) - - J_right = jacobian(right_ids, right_mask) - J_left = jacobian(left_ids, left_mask) - assert torch.allclose(J_left, J_right, atol=1e-4) - - -def test_capture_h_l_last_layer_is_pre_norm_residual(): - """transformers >= 5.x ties hidden_states[-1] to the post-norm residual; h_l must not be.""" - model, tokenizer = load_model("sshleifer/tiny-gpt2", torch.device("cpu"), torch.float32) - input_ids = tokenizer(["The cat sat."], return_tensors="pt")["input_ids"] - attention_mask = torch.ones_like(input_ids) - last = layer_indices(model)[-1] - base = _base_model(model) - - block_out: list[torch.Tensor] = [] - - def capture_block_output(module, input_tuple, output): - block_out.append((output[0] if isinstance(output, tuple) else output).detach().clone()) - - handle = base.h[last].register_forward_hook(capture_block_output) - try: - base( - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=get_position_ids(attention_mask), - return_dict=True, - ) - finally: - handle.remove() - - h_l = _capture_h_l(model, input_ids, attention_mask, last) - assert torch.allclose(h_l, block_out[0], atol=1e-6) - assert not torch.allclose(h_l, base.ln_f(block_out[0]), atol=1e-4) - - # Train-path preservation: with target_layer=None the last layer maps to itself. - z = _run_from_layer(model, last, h_l, input_ids, attention_mask, target_layer=None) - assert torch.equal(z, h_l) - - -def test_target_layer_none_runs_single_forward(): - """target_layer=None must not re-run the model (one embedding forward per pass).""" - model, tokenizer = load_model("sshleifer/tiny-gpt2", torch.device("cpu"), torch.float32) - input_ids = tokenizer(["The cat sat."], return_tensors="pt")["input_ids"] - attention_mask = torch.ones_like(input_ids) - wte = _base_model(model).wte - calls = {"n": 0} - handle = wte.register_forward_hook( - lambda module, inp, out: calls.__setitem__("n", calls["n"] + 1) - ) - try: - h_l = _capture_h_l(model, input_ids, attention_mask, layer_idx=0) - calls["n"] = 0 - _run_from_layer(model, 0, h_l, input_ids, attention_mask, target_layer=None) - assert calls["n"] == 1 - - calls["n"] = 0 - train_jacobian_lens( - model, - input_ids, - target_layer=None, - dtype=torch.float32, - batch_size=1, - output_dim_chunk=16, - ) - # capture(l0) + run(l0) + capture(l1); l1 short-circuits before any forward. - assert calls["n"] == 3 - finally: - handle.remove() - - -def test_vjp_restores_requires_grad_flags(): - model, tokenizer = load_model("sshleifer/tiny-gpt2", torch.device("cpu"), torch.float32) - input_ids = tokenizer(["The cat sat."], return_tensors="pt")["input_ids"] - attention_mask = torch.ones_like(input_ids) - params = list(model.parameters()) - for i, p in enumerate(params): - p.requires_grad_(i % 2 == 0) - saved = [p.requires_grad for p in params] - try: - _average_jacobian_for_layer( - model, input_ids, attention_mask, layer_idx=0, output_dim_chunk=16 - ) - assert [p.requires_grad for p in params] == saved - finally: - for p in params: - p.requires_grad_(True) - - -def test_frozen_qk_detaches_gpt2_c_attn_qk(): - """Gradients must not flow into c_attn's Q/K weight columns; V columns must still get grad.""" - model, _ = load_model("sshleifer/tiny-gpt2", torch.device("cpu"), torch.float32) - n_embd = model.config.n_embd - c_attn = model.transformer.h[0].attn.c_attn - hidden = torch.randn(1, 4, n_embd, requires_grad=True) - handles = _attach_frozen_qk_hooks(model) - try: - out = c_attn(hidden) - out.backward(torch.ones_like(out)) - finally: - for h in handles: - h.remove() - grad = c_attn.weight.grad - assert grad is not None - # Conv1D weight is (nx, nf); columns [0:n]=Q, [n:2n]=K, [2n:3n]=V. - assert grad[:, : 2 * n_embd].abs().max().item() == 0.0 - assert grad[:, 2 * n_embd :].abs().max().item() > 0.0 - - 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))