diff --git a/.gitignore b/.gitignore index 81900be..5a4f8b1 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ __pycache__/ artifacts/*.jac artifacts/*.ckpt uv.lock +docs/reviews/*.raw.jsonl diff --git a/docs/reviews/code.md b/docs/reviews/code.md new file mode 100644 index 0000000..559c2cc --- /dev/null +++ b/docs/reviews/code.md @@ -0,0 +1,56 @@ +Now let me verify one more file — checking if there are tests or any additional files:Here is the full review. + +--- + +## Code Review: jsteer — Jacobian pullback steering library + +### Summary +This is a thin, well-structured wrapper around `jlens` (Jacobian fitting/caching) and `steering-lite` (hook runtime). Fit once, derive steering vectors nearly for free. The math is parity-gated (cos > 0.999 vs direct VJP), so the review focuses on wiring, device/dtype handling, silent-failure anti-patterns, and API misuse. Overall the code is clean and follows the project's "fail fast, no defensive programming" convention, but has a few sharp edges. + +### Important (should fix) + +- **`jsteer/vjp.py:54`** — `assert max(layers) < target_layer` is an assert, not a `ValueError`. Asserts are stripped when Python runs with `-O`, which would let through invalid layer configs and produce incorrect gradients (source layers at or past the target layer). Replace with: + ```python + if max(layers) >= target_layer: + raise ValueError(f"source layers {layers} must be < {target_layer}; got max={max(layers)}") + ``` + The convention of raising `ValueError` for this is established in `Jacobian.pullback` which does exactly that for missing layers. + +- **`jsteer/jacobian.py:92-105`** (`_h_bar_final`) — No guard for empty `prompts`. If called with an empty list, `acc` stays `None` and `acc / n` raises `TypeError` (None / 0) rather than a clear error. Given that `persona_vector` and `persona_topk_vector` both call this with caller-supplied prompt lists, a well-meaning empty-pass is plausible. Add a check at the top: + ```python + if not prompts: + raise ValueError("prompts must not be empty") + ``` + +### Suggestions + +- **`jsteer/jacobian.py:60-68`** (`_steer_layers`) — `tuple(sorted(int(l) for l in layers))` silently truncates float bands (e.g. `(0.5, 0.8)` → `(0, 0)`) rather than rejecting them. Float bands are only meaningful at fit time (`_resolve_layers` handles them), and `_steer_layers` is post-fit. A `ValueError` for float inputs would make the contract explicit and prevent a user from accidentally passing `layers=(0.5, 0.8)` and getting nonsense layers (0, 0). + +- **`jsteer/jacobian.py:166-179`** (`persona_topk_vector`) — Calls `from_hf(model, tok)` on its first line, then `_h_bar_final` twice (which internally also calls `from_hf`). Three redundant `HFLensModel` constructions: three iterations over all params to freeze, three layout detections. Each `from_hf` is ~1-2 ms plus linear in param count, so for the 0.6B model it's invisible but for larger models it adds up. Extract `lm = from_hf(model, tok)` once and pass `n_layers` to `_h_bar_final` (or refactor `_h_bar_final` to accept a pre-made `lm`). + +- **`jsteer/jacobian.py:96-98`** (`_h_bar_final`) — Uses `model(**enc)` (the full HF model including LM head forward) rather than `lm.forward(input_ids)` (residual stack only). The LM head computation is wasted work done for every batch. Under `@torch.no_grad()` the overhead is minor but inconsistent: `pullback_vjp` uses `model(**enc)` too (needs the full model because of hook placement), but `jlens.fit` correctly uses `lm.forward`. Would be cleaner to use `lm.forward` here since only residuals are needed. + +- **`jsteer/jacobian.py:171,185-186`** (`persona_topk_vector`) — `cots[name] = W_U[top.indices].float().mean(0).cpu()` reads raw unembedding rows (no final norm), while `lm.unembed(...)` above goes through final norm to pick the top-k tokens. This is intentional (consistent with `_word_cotangent`'s raw-row convention and the docstring), but it does mean the "most evoked tokens" are selected via the full logit pipeline while the downstream cotangent uses the raw dueling basis. A single-line comment explaining the asymmetry would help future readers. + +- **`jsteer/vjp.py:29-34`** (`_valid_mask`) — `mask & attention_mask.bool()` redundantly masks with the attention mask after already filtering by `pos < real_len - 1`. For standard HF right-padded batches these are equivalent, but the redundancy isn't harmful. Fine to leave, but a one-line comment that it's a belt-and-suspenders check would prevent a future reader from "simplifying" it and breaking left-padded or non-square attention mask scenarios. + +- **`jsteer/jacobian.py:188`** (`random_vector`) — Generates directions on CPU without an explicit `dtype` argument. `torch.randn` defaults to `torch.float32`, which is correct. If this ever needs to match the model dtype (e.g. bf16), it would need updating. + +- **`jsteer/applies.py`** — The `_extract_stub` and registration loop are clean but the docstring in `_extract_stub` could mention that `steering_lite.train` is the entry point being blocked. Currently the error message explains what to do, but a developer seeing "NotImplementedError: jsteer methods are extracted via Jacobian..." from inside `steering_lite.train()` might not immediately connect the dots. Minor. + +### Positive + +- **`jsteer/jacobian.py:109-117`** (`pullback`) — Pre-validates cotangent shape and layer membership with clear `ValueError` messages before touching tensors. Exactly the right fail-fast pattern. + +- **`jsteer/vjp.py:62-64`** — The zero-valid-positions check catches short prompts early with a clear error, preventing silent zeros downstream. + +- **`jsteer/applies.py`** — `apply_add_last` correctly degrades to `apply_add` when span ≥ sequence length (the slicing `y[:, :-k, :]` yields empty, `cat` reconstructs the full sequence). Documented and correct. + +- **Sign convention consistency** — All three concept-method docstrings explicitly state what `+C` does, and the pullback computation (`w @ J_l`, i.e., `J_l^T @ w`) is consistent: `+C` enhances the named concept. + +- **`_to_vector` layout** — The `stacked["v"].unsqueeze(0)` with `k=1` leading dim matches `steering_lite`'s `mean_diff` layout byte-for-byte, so calibration and serialization reuse the upstream code unchanged. This is the correct integration pattern. + +### Verdict +**APPROVE** with minor fixes. + +The `assert` → `ValueError` in `vjp.py:54` and the empty-prompts guard in `_h_bar_final` are the two changes worth making before shipping. Everything else is suggestions. The wiring is correct, the sign conventions are consistent, the jlens API is used properly (no reinvention of the estimator), and the parity gate confirms numerical equivalence. \ No newline at end of file diff --git a/jsteer/jacobian.py b/jsteer/jacobian.py index 519886c..99f0b02 100644 --- a/jsteer/jacobian.py +++ b/jsteer/jacobian.py @@ -195,6 +195,11 @@ class Jacobian: a fit-time concept -- the lens doesn't know n_layers to resolve them).""" if layers is None: return tuple(self.lens.source_layers) + if any(isinstance(l, float) for l in layers): + # Claude: int() would silently truncate (0.5, 0.8) -> layer 0 and steer + # the wrong layer; float bands only exist at fit time (external review). + raise ValueError(f"float layer bands are fit-time only; got {layers}, " + f"pass explicit ints from .layers={self.layers}") return tuple(sorted(int(l) for l in layers)) # -- concept -> vector ------------------------------------------------------- @@ -237,6 +242,9 @@ class Jacobian: toks = [tok.decode([i]) for i in top.indices.tolist()] logger.info(f"persona_topk {name} top-{k}: {toks}") # read your data: # gibberish/punctuation here means the persona mean is off-manifold + # Claude: asymmetry is intentional -- token SELECTION goes through the + # full logit pipeline (final norm) above, but the cotangent uses raw + # W_U rows to match _word_cotangent's verified convention. cots[name] = W_U[top.indices].float().mean(0).cpu() cfg = JacobianPersonaTopkC(layers=self._steer_layers(layers)) return self.pullback(cots["pos"] - cots["neg"], cfg) diff --git a/jsteer/vjp.py b/jsteer/vjp.py index 21961c0..a3a77ef 100644 --- a/jsteer/vjp.py +++ b/jsteer/vjp.py @@ -38,6 +38,8 @@ def _valid_mask(attention_mask: Tensor, skip_first: int) -> Tensor: real_len = attention_mask.sum(dim=1, keepdim=True) # [B, 1] pos = torch.arange(attention_mask.shape[1], device=attention_mask.device) mask = (pos[None, :] >= skip_first) & (pos[None, :] < real_len - 1) + # Claude: the & with attention_mask is redundant for right-padded batches + # (pos < real_len-1 already excludes pads) but guards non-right-padded input. return mask & attention_mask.bool() @@ -49,7 +51,8 @@ def pullback_vjp(model, tok, prompts: list[str], layers, cotangent: Tensor, *, lm = from_hf(model, tok) # freezes params, locates blocks; grads flow to target_layer = lm.n_layers - 1 # activations only layers = _resolve_layers(layers, lm.n_layers) - assert max(layers) < target_layer, f"source layers {layers} must be < {target_layer}" + if max(layers) >= target_layer: + raise ValueError(f"source layers {layers} must be < target {target_layer}") d = cotangent.shape[0] G = {l: torch.zeros(d, dtype=torch.float32, device=model.device) for l in layers} count = 0