core algo: Jacobian fit/cache/pullback (jlens) + word/persona vectors (steering-lite runtime)

Ported from the verified j-steer-dev experiment (word pullback beat random
on 3/5 moral foundations, Qwen3-4B n=3). jacobian.py wraps jlens fit/save/
load and derives steering vectors by CPU matvec against the cached J;
vjp.py is the one-backward parity path; applies.py registers the methods
and delivery modes into steering-lite.

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>

.gitignore

@@ -0,0 +1,6 @@
+.venv/
+__pycache__/
+*.egg-info/
+artifacts/*.jac
+artifacts/*.ckpt
+uv.lock

AGENTS.md

@@ -0,0 +1,43 @@
+# jsteer -- agent notes
+
+Plan of record: /home/wassname/.claude/plans/review-specs-00-minimal-experiment-md-an-peppy-sky.md
+Evidence base: ../j-steer-dev/docs/RESEARCH_JOURNAL.md (verified 3/5 word-steering result).
+
+## What this is
+
+repeng-style UX for Jacobian pullback steering. Core algo files are WRITTEN
+(by the main agent, ported from verified j-steer-dev code -- do not rewrite
+the math, it is parity-gated against the verified experiment):
+
+- `jsteer/jacobian.py` -- Jacobian.fit/save/load (wraps jlens, the
+  researchers' verified primary code) + word/persona/persona_topk/random
+  vectors -> steering_lite.Vector.
+- `jsteer/applies.py` -- steering-lite method registration + delivery modes.
+- `jsteer/vjp.py` -- direct VJP path; the parity reference for the cache.
+
+Runtime is steering-lite: `with v(model, C=8): model.generate(...)`.
+
+## Remaining work (task list has details; U-numbers from the plan)
+
+1. uv scaffold: `uv sync` (torch cu121+ index if needed), fix any import errors
+   fail-fast (no defensive fallbacks). LOCAL DEV: you may switch
+   [tool.uv.sources] to path deps (../j-steer-dev/docs/vendor/jacobian-lens
+   and ../../lite/steering-lite, editable) if the git fetches are slow --
+   leave a comment saying which is active and why.
+2. Smoke on Qwen/Qwen3-0.6B: tiny fit (8 short web-text prompts, mid layers,
+   dim_batch 8), word_vector(["happy","joy"]), generate at C in {-8, 0, 8},
+   print FULL first prompt + generations (token-efficient-logging skill).
+3. U1 parity gate BEFORE demos: cos(Jacobian-cache pullback, word_vector_vjp)
+   per layer > 0.999, same prompts/max_length/skip_first. If it fails, that is
+   a bug in the wiring (the math is linear-identical), debug do not tune.
+4. 0.6B real fit (~64 prompts) cached to artifacts/; 4B via pueue (label
+   why:/resolve:).
+5. notebooks/word_steering.ipynb (hello-world), persona_steering.ipynb
+   (persona variants are EXPERIMENTAL -- they failed specificity controls in
+   j-steer-dev; keep that framing), lens_readout.ipynb optional.
+6. README: classic-repeng length, honest evidence section.
+
+## Style
+
+Fail fast, no defensive programming, loguru, no LLM-tell prose in README.
+Comments marked as Claude-authored where opinionated.

jsteer/__init__.py

@@ -0,0 +1,13 @@
+"""jsteer: fit a model's full Jacobian once, then steer any word or persona.
+
+    from jsteer import Jacobian
+    jac = Jacobian.fit(model, tok, prompts)         # expensive, once, cacheable
+    v = jac.word_vector(model, tok, ["authority"])  # instant matvec
+    with v(model, C=8):                             # steering-lite runtime
+        model.generate(**inputs)
+"""
+from . import applies  # noqa: F401  -- registers methods into steering-lite's REGISTRY
+from .jacobian import Jacobian
+from .vjp import pullback_vjp, word_vector_vjp
+
+__all__ = ["Jacobian", "pullback_vjp", "word_vector_vjp"]

jsteer/applies.py

@@ -0,0 +1,159 @@
+"""steering-lite method registration + modular delivery of jacobian vectors.
+
+(drafted by Claude, ported from the verified j-steer-dev experiment code)
+
+jsteer vectors are plain `steering_lite.Vector` objects: one unit direction v
+per layer in `stacked["v"]` with a leading k=1 dim `[1, d]` (byte-identical
+layout to steering-lite's mean_diff), so attach / calibrate / save / `with
+v(model, C=...)` all work unchanged.
+
+Extraction never goes through steering-lite's `train()` (it needs gradients
+that train's no_grad path can't give) -- it lives in `jacobian.py` (cached
+full-J pullback) and `vjp.py` (direct VJP). The `extract` entries here are
+stubs that say so.
+
+DELIVERY of v to the residual stream is decoupled from extraction: the same v
+can be added everywhere, gated to the last positions, or overwrite a span.
+`cfg.apply_mode` selects the mode; adding a mode = one function + one
+APPLY_REGISTRY entry.
+
+Protocol (steering-lite config.py Method.apply):
+    apply(mod, x, y, shared, stacked, cfg) -> y_new   # same shape [b, s, d]
+
+Position semantics: generation uses LEFT padding, so the last real token is
+position -1; `cfg.apply_span` selects how many trailing positions to target.
+Note that during incremental generation (KV cache) every decode step has s=1,
+so add_last touches each generated token but only the tail of the prefill.
+
+Sign conventions (set by the cotangent in jacobian.py):
+    jacobian_word          +C raises the words' output logits
+    jacobian_persona       +C moves toward the POS persona
+    jacobian_persona_topk  +C moves toward the POS persona's evoked vocabulary
+    random                 norm-matched control, no meaning
+"""
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor
+
+from steering_lite.config import SteeringConfig, register, register_config
+
+ε = 1e-8
+
+
+# --- configs -----------------------------------------------------------------
+# One config per method name so saved vectors deserialize with provenance.
+# All share the same knobs: apply_mode (delivery) + apply_span (tail width).
+
+@register_config
+@dataclass
+class JacobianWordC(SteeringConfig):
+    method: str = "jacobian_word"
+    normalize: bool = True
+    apply_mode: str = "add"
+    apply_span: int = 1
+
+
+@register_config
+@dataclass
+class JacobianPersonaC(SteeringConfig):
+    method: str = "jacobian_persona"
+    normalize: bool = True
+    apply_mode: str = "add"
+    apply_span: int = 1
+
+
+@register_config
+@dataclass
+class JacobianPersonaTopkC(SteeringConfig):
+    method: str = "jacobian_persona_topk"
+    normalize: bool = True
+    apply_mode: str = "add"
+    apply_span: int = 1
+
+
+@register_config
+@dataclass
+class RandomC(SteeringConfig):
+    method: str = "random"
+    normalize: bool = True
+    apply_mode: str = "add"
+    apply_span: int = 1
+
+
+# --- delivery modes ----------------------------------------------------------
+
+def _v_sum(stacked: dict[str, Tensor], y: Tensor) -> Tensor:
+    """The per-layer direction v, summed over the k-stack, on y's device/dtype."""
+    return stacked["v"].to(y).sum(dim=0)  # [d]
+
+
+def apply_add(mod, x, y, shared, stacked, cfg) -> Tensor:
+    """y += coeff * v at ALL positions (the verified default; same delivery as
+    steering-lite's mean_diff, so calibrated coeffs are comparable)."""
+    v = _v_sum(stacked, y)
+    return y + cfg.coeff * v
+
+
+def apply_add_last(mod, x, y, shared, stacked, cfg) -> Tensor:
+    """y[:, -k:] += coeff * v -- nudge only the last k positions (decision
+    region). k = cfg.apply_span; k >= s degenerates to apply_add."""
+    v = _v_sum(stacked, y)
+    k = cfg.apply_span
+    head = y[:, :-k, :]
+    tail = y[:, -k:, :] + cfg.coeff * v
+    return torch.cat([head, tail], dim=1)
+
+
+def apply_replace_last(mod, x, y, shared, stacked, cfg) -> Tensor:
+    """Overwrite the last k positions with the concept direction at each
+    position's original magnitude: energy from y, direction from v, strength
+    from coeff. A "virtual token" injection that keeps the [b, s, d] shape
+    (true sequence insertion would break RoPE / attention mask / KV cache)."""
+    v = _v_sum(stacked, y)
+    v_unit = v / (v.norm() + ε)
+    k = cfg.apply_span
+    tail = y[:, -k:, :]                           # [b, k, d]
+    energy = tail.norm(dim=-1, keepdim=True)      # [b, k, 1]
+    new_tail = energy * (cfg.coeff * v_unit)
+    head = y[:, :-k, :]
+    return torch.cat([head, new_tail], dim=1)
+
+
+APPLY_REGISTRY: dict[str, Callable[..., Tensor]] = {
+    "add": apply_add,
+    "add_last": apply_add_last,
+    "replace_last": apply_replace_last,
+}
+
+
+def apply_dispatch(mod, x, y, shared, stacked, cfg) -> Tensor:
+    fn = APPLY_REGISTRY.get(cfg.apply_mode)
+    if fn is None:
+        raise KeyError(
+            f"unknown apply_mode={cfg.apply_mode!r}; registered: {list(APPLY_REGISTRY)}")
+    return fn(mod, x, y, shared, stacked, cfg)
+
+
+# --- method registration -----------------------------------------------------
+
+def _extract_stub(pos_acts, neg_acts, cfg):
+    raise NotImplementedError(
+        "jsteer methods are extracted via Jacobian.{word,persona,persona_topk}_vector "
+        "(jacobian.py) or vjp.py, not steering_lite.train -- they need gradients "
+        "that train's no_grad path can't provide.")
+
+
+def _register_method(method_name: str) -> None:
+    @register
+    class _M:  # noqa: N801 -- registry keys on .name, class name irrelevant
+        name = method_name
+        extract = staticmethod(_extract_stub)
+        apply = staticmethod(apply_dispatch)
+
+
+for _name in ("jacobian_word", "jacobian_persona", "jacobian_persona_topk", "random"):
+    _register_method(_name)

jsteer/jacobian.py

@@ -0,0 +1,264 @@
+"""Full-Jacobian extraction for steering: fit once, steer any concept.
+
+(drafted by Claude, ported from the verified j-steer-dev experiment code)
+
+The method (verified in j-steer-dev: word steering beat a norm-matched random
+control on 3/5 moral foundations, Qwen3-4B, n=3 seeds):
+
+    v_l = unit( J_l^T @ w )
+
+where `J_l = E_prompts[ d h_final / d h_l ]` is the jlens position-averaged
+Jacobian -- the researchers' verified estimator, reused via `jlens.fitting.fit`
+(never reimplemented) -- and `w` is a cotangent in the FINAL-layer basis naming
+the concept to steer:
+
+    word_vector           w = mean unembedding row of the words     VERIFIED
+    persona_vector        w = h_bar(pos) - h_bar(neg)               EXPERIMENTAL*
+    persona_topk_vector   w = contrast of the top-k tokens each     EXPERIMENTAL
+                          persona evokes at the final layer
+
+    * persona-contrast pullbacks FAILED specificity controls in j-steer-dev
+      (moved the target axis no more than an unrelated persona did). Shipped
+      for experimentation, not as a recommendation.
+
+Why cache the full J: by linearity `mean_p(J_p)^T w = mean_p(J_p^T w)`, so a
+vector pulled back through the cached pooled Jacobian is numerically the same
+vector the direct per-prompt VJP produces (see vjp.py; parity-tested). Fitting
+is the expensive step (one forward + ceil(d_model/dim_batch) backwards per
+prompt); afterwards every concept vector is a CPU matvec.
+
+The Jacobian is always fit against the FINAL layer basis (jlens default), so
+cotangents are measured there: unembedding rows live there natively, persona
+means are recorded there.
+
+Vectors come out as `steering_lite.Vector` (unit direction per layer in
+stacked["v"], k=1 leading dim -- mean_diff's exact layout), so steering is:
+
+    with v(model, C=8):
+        model.generate(**inputs)
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from jlens.fitting import fit as _jlens_fit
+from jlens.hf import HFLensModel, from_hf
+from jlens.hooks import ActivationRecorder
+from jlens.lens import JacobianLens
+from loguru import logger
+from torch import Tensor
+from tqdm.auto import tqdm
+
+from steering_lite.config import SteeringConfig
+from steering_lite.vector import Vector
+
+from .applies import (
+    JacobianPersonaC,
+    JacobianPersonaTopkC,
+    JacobianWordC,
+    RandomC,
+    ε,
+)
+
+
+# --- small helpers ------------------------------------------------------------
+
+def _unit(v: Tensor) -> Tensor:
+    return v / (v.norm() + ε)
+
+
+def _resolve_layers(layers, n_layers: int) -> list[int]:
+    """None -> all layers below final; (lo, hi) floats -> fraction band;
+    iterable of ints -> as-is (negatives count from the end)."""
+    if layers is None:
+        return list(range(n_layers - 1))
+    layers = tuple(layers)
+    if len(layers) == 2 and all(isinstance(x, float) for x in layers):
+        lo, hi = (int(layers[0] * n_layers), int(layers[1] * n_layers))
+        return list(range(lo, min(hi, n_layers - 1)))
+    return sorted({l + n_layers if l < 0 else l for l in layers})
+
+
+def _to_vector(cfg: SteeringConfig, per_layer: dict[int, Tensor]) -> Vector:
+    """Wrap unit directions as a steering-lite Vector (mean_diff's layout:
+    stacked["v"] with leading k=1 dim, shared empty)."""
+    shared = {l: {} for l in per_layer}
+    stacked = {l: {"v": _unit(v.float()).unsqueeze(0)} for l, v in per_layer.items()}
+    return Vector(cfg, shared, stacked)
+
+
+def _word_cotangent(model, tok, words: list[str]) -> Tensor:
+    """Mean unembedding row over `words` (first sub-token of each): the
+    final-basis direction that most raises those tokens' output logits.
+    Pulling THIS back through J^T is the pure concept->residual map -- no
+    persona pairs, so no persona-bundle confound. +C enhances the concept."""
+    W_U = model.lm_head.weight                                   # [vocab, d]
+    ids = [tok(w, add_special_tokens=False).input_ids[0] for w in words]
+    rows = W_U[torch.tensor(ids, device=W_U.device)].float()     # [n, d]
+    cot = rows.mean(0)
+    logger.info(f"word cotangent: {words} -> first-subtoken ids={ids} |w|={cot.norm():.3f}")
+    return cot
+
+
+@torch.no_grad()
+def _h_bar_final(model, tok, prompts: list[str], *, batch_size: int = 8,
+                 max_length: int = 384, label: str = "") -> Tensor:
+    """Mean last-token activation at the FINAL layer over `prompts`
+    (right-padded batches; last real token located via the attention mask)."""
+    lm = from_hf(model, tok)                                     # layout detection only
+    target_layer = lm.n_layers - 1
+    acc, n = None, 0
+    for i in tqdm(range(0, len(prompts), batch_size), desc=f"h_bar {label}",
+                  mininterval=30, maxinterval=60):
+        batch = prompts[i:i + batch_size]
+        enc = tok(batch, return_tensors="pt", padding=True, truncation=True,
+                  max_length=max_length, padding_side="right").to(model.device)
+        with ActivationRecorder(lm.layers, at=[target_layer]) as rec:
+            model(**enc)
+        act = rec.activations[target_layer]                      # [B, S, d]
+        last_idx = enc["attention_mask"].sum(dim=1) - 1          # [B]
+        last = act[torch.arange(act.shape[0]), last_idx].float()
+        acc = last.sum(0) if acc is None else acc + last.sum(0)
+        n += last.shape[0]
+    return acc / n
+
+
+# --- the core object ----------------------------------------------------------
+
+@dataclass
+class Jacobian:
+    """A fitted, cached, model-specific bundle of per-layer Jacobians
+    `{layer: J_l [d, d]}` (a `jlens.JacobianLens`) plus the concept->vector
+    methods. Fit once (expensive), derive steering vectors forever (matvec)."""
+
+    lens: JacobianLens
+
+    # -- fit / persist (all delegated to the researchers' jlens code) ----------
+
+    @classmethod
+    def fit(cls, model, tok, prompts: list[str], *, layers=None, dim_batch: int = 8,
+            max_seq_len: int = 128, checkpoint_path: str | None = None,
+            compile: bool = False) -> "Jacobian":
+        """Fit `J_l` on `prompts` (generic text; jlens guidance is ~100+ for
+        lens-quality pooling). Cost: 1 forward + ceil(d_model/dim_batch)
+        backwards per prompt -- ALL source layers come from the same backwards,
+        so fitting more layers costs memory, not compute. `checkpoint_path`
+        makes the fit resumable (atomic writes)."""
+        lm = from_hf(model, tok, compile=compile)
+        source_layers = _resolve_layers(layers, lm.n_layers)
+        lens = _jlens_fit(lm, prompts, source_layers=source_layers,
+                          dim_batch=dim_batch, max_seq_len=max_seq_len,
+                          checkpoint_path=checkpoint_path)
+        return cls(lens=lens)
+
+    def save(self, path: str) -> None:
+        self.lens.save(path)      # fp16 by default; jlens-compatible file
+
+    @classmethod
+    def load(cls, path: str) -> "Jacobian":
+        return cls(lens=JacobianLens.load(path))
+
+    @classmethod
+    def from_pretrained(cls, name_or_path: str, **kw) -> "Jacobian":
+        """Local file/dir or HuggingFace Hub repo_id (see JacobianLens)."""
+        return cls(lens=JacobianLens.from_pretrained(name_or_path, **kw))
+
+    @property
+    def layers(self) -> list[int]:
+        return self.lens.source_layers
+
+    def __repr__(self) -> str:
+        return f"Jacobian({self.lens!r})"
+
+    # -- the one pullback -------------------------------------------------------
+
+    def pullback(self, cotangent: Tensor, cfg: SteeringConfig) -> Vector:
+        """v_l = unit(J_l^T @ w) for every layer in cfg.layers.
+
+        J_l rows are output dims (each row is the gradient of one final-basis
+        dim over the source layer), so the pullback is `w @ J_l`. Computed on
+        CPU fp32 against the cached matrices -- no model, no backward."""
+        w = cotangent.detach().float().cpu()
+        if w.shape != (self.lens.d_model,):
+            raise ValueError(f"cotangent shape {tuple(w.shape)} != ({self.lens.d_model},)")
+        missing = set(cfg.layers) - set(self.lens.source_layers)
+        if missing:
+            raise ValueError(f"layers {sorted(missing)} not fitted; have {self.layers}")
+        per_layer = {l: w @ self.lens.jacobians[l] for l in cfg.layers}
+        logger.info(f"{cfg.method} per-layer |J^T w| (pre-norm): " +
+                    " ".join(f"{l}:{per_layer[l].norm():.3g}" for l in cfg.layers))
+        return _to_vector(cfg, per_layer)
+
+    def _steer_layers(self, layers) -> tuple[int, ...]:
+        """None -> all fitted layers; else explicit int indices (float bands are
+        a fit-time concept -- the lens doesn't know n_layers to resolve them)."""
+        if layers is None:
+            return tuple(self.lens.source_layers)
+        return tuple(sorted(int(l) for l in layers))
+
+    # -- concept -> vector -------------------------------------------------------
+
+    def word_vector(self, model, tok, words: list[str], *, layers=None) -> Vector:
+        """VERIFIED method: pull the words' unembedding direction back through
+        the Jacobian. +C enhances the concept, -C suppresses it."""
+        cfg = JacobianWordC(layers=self._steer_layers(layers))
+        return self.pullback(_word_cotangent(model, tok, words), cfg)
+
+    def persona_vector(self, model, tok, pos_prompts: list[str],
+                       neg_prompts: list[str], *, layers=None,
+                       batch_size: int = 8) -> Vector:
+        """EXPERIMENTAL: pull the persona activation contrast back through the
+        Jacobian. Persona-contrast pullbacks failed specificity controls in
+        j-steer-dev -- prefer word_vector for targeted steering."""
+        h_pos = _h_bar_final(model, tok, pos_prompts, batch_size=batch_size, label="pos")
+        h_neg = _h_bar_final(model, tok, neg_prompts, batch_size=batch_size, label="neg")
+        logger.info(f"h_bar_diff |pos|={h_pos.norm():.3f} |neg|={h_neg.norm():.3f} "
+                    f"|diff|={ (h_pos - h_neg).norm():.3f}")
+        cfg = JacobianPersonaC(layers=self._steer_layers(layers))
+        return self.pullback(h_pos - h_neg, cfg)
+
+    def persona_topk_vector(self, model, tok, pos_prompts: list[str],
+                            neg_prompts: list[str], *, k: int = 8, layers=None,
+                            batch_size: int = 8) -> Vector:
+        """EXPERIMENTAL: persona -> vocabulary bottleneck -> word pullback.
+        Read each persona's final-layer mean through the unembedding, take the
+        top-k tokens it most evokes, contrast the two token sets' unembedding
+        rows, pull that back. Composes the persona signal with the verified
+        word mechanism; untested for specificity."""
+        lm = from_hf(model, tok)
+        h_pos = _h_bar_final(model, tok, pos_prompts, batch_size=batch_size, label="pos")
+        h_neg = _h_bar_final(model, tok, neg_prompts, batch_size=batch_size, label="neg")
+        W_U = model.lm_head.weight                                # [vocab, d]
+        cots = {}
+        for name, h in (("pos", h_pos), ("neg", h_neg)):
+            logits = lm.unembed(h.to(model.device).to(model.dtype)).float()
+            top = logits.topk(k)
+            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
+            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)
+
+    def random_vector(self, *, seed: int = 0, layers=None) -> Vector:
+        """Norm-matched control: unit random direction per layer. Any honest
+        demo/eval should show the concept vector beating THIS at the same C."""
+        gen = torch.Generator().manual_seed(seed)
+        cfg = RandomC(layers=self._steer_layers(layers), seed=seed)
+        per_layer = {l: torch.randn(self.lens.d_model, generator=gen)
+                     for l in cfg.layers}
+        return _to_vector(cfg, per_layer)
+
+    # -- bonus: the lens's native forward readout --------------------------------
+
+    def lens_topk(self, model, tok, prompt: str, layer: int, *, k: int = 10,
+                  position: int = -1) -> list[tuple[str, float]]:
+        """What the model 'thinks' at `layer`: transport the residual to the
+        final basis with J_l and decode. jlens's native use, handy in demos."""
+        lm = from_hf(model, tok)
+        lens_logits, _, _ = self.lens.apply(lm, prompt, layers=[layer],
+                                            positions=[position])
+        top = lens_logits[layer][0].topk(k)
+        return [(tok.decode([i]), float(v)) for i, v in
+                zip(top.indices.tolist(), top.values.tolist())]

jsteer/vjp.py

@@ -0,0 +1,93 @@
+"""Direct VJP pullback: one concept without paying for the full Jacobian.
+
+(drafted by Claude, ported from the verified j-steer-dev experiment code)
+
+Computes the SAME vector as `Jacobian.pullback` -- by linearity
+`mean_p(J_p^T w) = mean_p(J_p)^T w` -- but contracts the cotangent inside the
+backward pass, so the cost is ONE backward per prompt instead of
+ceil(d_model/dim_batch). Use this when you want a single concept and don't
+need the reusable cache; use it in tests as the parity reference for the
+cached path (cos > 0.999 per layer expected, fp16 storage being the only gap).
+
+Estimator conventions are jlens's exactly (this is the code path that produced
+the verified j-steer-dev result): cotangent placed at every valid target
+position of the final layer, gradient read at every valid source position and
+meaned, positions before skip_first=16 excluded (attention sinks), final
+position excluded (no next-token target). Right-padded batches: the valid mask
+excludes pads, so batching does not change the estimate.
+"""
+from __future__ import annotations
+
+import torch
+from jlens.fitting import SKIP_FIRST_N_POSITIONS
+from jlens.hf import from_hf
+from jlens.hooks import ActivationRecorder
+from loguru import logger
+from torch import Tensor
+from tqdm.auto import tqdm
+
+from steering_lite.vector import Vector
+
+from .applies import JacobianWordC
+from .jacobian import _resolve_layers, _to_vector, _word_cotangent
+
+
+def _valid_mask(attention_mask: Tensor, skip_first: int) -> Tensor:
+    """Boolean [B, S]: right-padded real tokens in [skip_first : real_len-1].
+    jlens valid_position_mask extended to a right-padded batch."""
+    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)
+    return mask & attention_mask.bool()
+
+
+def pullback_vjp(model, tok, prompts: list[str], layers, cotangent: Tensor, *,
+                 batch_size: int = 8, max_length: int = 128,
+                 skip_first: int = SKIP_FIRST_N_POSITIONS) -> dict[int, Tensor]:
+    """Per-layer mean over `prompts` of J_l^T @ cotangent, via one backward per
+    batch (grads for every source layer come from the same backward)."""
+    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}"
+    d = cotangent.shape[0]
+    G = {l: torch.zeros(d, dtype=torch.float32, device=model.device) for l in layers}
+    count = 0
+    for i in tqdm(range(0, len(prompts), batch_size), desc="pullback_vjp",
+                  mininterval=30, maxinterval=60):
+        batch = prompts[i:i + batch_size]
+        enc = tok(batch, return_tensors="pt", padding=True, truncation=True,
+                  max_length=max_length, padding_side="right").to(model.device)
+        valid = _valid_mask(enc["attention_mask"], skip_first)    # [B, S] bool
+        if valid.sum(dim=1).min() == 0:
+            raise ValueError(f"a prompt has 0 valid positions "
+                             f"(too short for skip_first={skip_first})")
+        with ActivationRecorder(lm.layers, at=[*layers, target_layer],
+                                start_graph_at=min(layers)) as rec, torch.enable_grad():
+            model(**enc)
+            h_final = rec.activations[target_layer]               # [B, S, d]
+            srcs = [rec.activations[l] for l in layers]
+            c = (cotangent.to(h_final.device).to(h_final.dtype).view(1, 1, d)
+                 * valid.unsqueeze(-1))
+            grads = torch.autograd.grad(h_final, srcs, grad_outputs=c)
+        den = valid.sum(dim=1, keepdim=True).float()              # [B, 1]
+        for l, g in zip(layers, grads):
+            v_b = (g.float() * valid.unsqueeze(-1)).sum(dim=1) / den   # [B, d]
+            G[l] += v_b.sum(0)
+        count += len(batch)
+    return {l: G[l] / count for l in layers}
+
+
+def word_vector_vjp(model, tok, prompts: list[str], words: list[str], *,
+                    layers=None, batch_size: int = 8, max_length: int = 128,
+                    skip_first: int = SKIP_FIRST_N_POSITIONS) -> Vector:
+    """The verified j-steer-dev method-0 extraction, self-contained: word
+    cotangent pulled back over `prompts` as linearization substrate. Same
+    vector as Jacobian.fit(model, tok, prompts).word_vector(...) when the
+    prompts, layers, skip_first and max length match."""
+    cot = _word_cotangent(model, tok, words)
+    G = pullback_vjp(model, tok, prompts, layers, cot, batch_size=batch_size,
+                     max_length=max_length, skip_first=skip_first)
+    logger.info("word_vector_vjp per-layer |v| (pre-norm): " +
+                " ".join(f"{l}:{v.norm():.3g}" for l, v in G.items()))
+    return _to_vector(JacobianWordC(layers=tuple(G)), G)

pyproject.toml

@@ -0,0 +1,31 @@
+[project]
+name = "jsteer"
+version = "0.1.0"
+description = "Fit a model's full Jacobian once (jlens), then steer any word or persona (steering-lite runtime). repeng-style UX for Jacobian pullback steering."
+readme = "README.md"
+requires-python = ">=3.11"
+license = "MIT"
+dependencies = [
+    "torch>=2.1",
+    "transformers>=4.51",
+    "accelerate>=1.6",
+    "jlens",           # the researchers' verified Jacobian estimator + cache format
+    "steering-lite",   # hook runtime: Vector, attach, `with v(model, C=...)`, calibrate
+    "loguru>=0.7",
+    "tqdm>=4.66",
+]
+
+[project.optional-dependencies]
+notebooks = ["matplotlib>=3.8", "ipykernel", "tabulate"]
+test = ["pytest"]
+
+[tool.uv.sources]
+jlens = { git = "https://github.com/anthropics/jacobian-lens" }
+steering-lite = { git = "https://github.com/wassname/steering-lite" }
+
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools]
+packages = ["jsteer"]