open-interpretability 0.0.20__py3-none-any.whl

openinterp/harness/moe.py

@@ -0,0 +1,134 @@
+"""The experts forward that openinterp runs mixture-of-experts models with (docs/moe.md).
+
+transformers 5 runs the experts of about 60 MoE architectures through
+``ExpertsInterface``: an experts class decorated with ``use_experts_implementation`` looks
+its forward up by name. :func:`experts_forward` is registered there as ``openinterp``,
+and :func:`use_openinterp_experts` selects it when a model loads. It is the ``eager``
+forward written out once for every layout the decorator describes, with two additions:
+
+* Each expert's hidden activation passes through the hook of :func:`expert_hook`
+  before the down projection. That is where an expert neuron is read and edited.
+* In a streamed model (``--memory stream``) the experts' tensors stay on disk, and the
+  forward reads only the routed experts' slices of them (:func:`expert_weight`). One
+  expert of Mistral Large 4 is about 113 MB in FP8, and all 128 of a layer 14.5 GB.
+
+``scripts/moe_experts_check.py`` and ``tests/test_moe.py`` compare it with ``eager``.
+"""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+
+import torch
+from torch import nn
+
+EXPERTS_IMPLEMENTATION = "openinterp"
+
+# fn(act [tokens, d_expert], expert index, token indices, routing weights) -> act
+ExpertHook = Callable[[torch.Tensor, int, torch.Tensor, torch.Tensor], torch.Tensor]
+
+_hooks = threading.local()  # per thread, as the layer-major review needs
+
+
+def _linear(x: torch.Tensor, weight: torch.Tensor, bias, transposed: bool) -> torch.Tensor:
+    weight = weight.to(x.device, x.dtype)  # a streamed slice is read to the CPU
+    y = x @ weight if transposed else x @ weight.T
+    return y if bias is None else y + bias.to(x.device, x.dtype)
+
+
+def experts_forward(self, hidden, top_k_index, top_k_weights) -> torch.Tensor:
+    """The routed experts' summed, weighted outputs ``[tokens, d_model]``.
+
+    ``top_k_index`` and ``top_k_weights`` are each token's chosen experts and their
+    routing weights ``[tokens, k]``. An index of ``num_experts`` or more is an
+    expert-parallel placeholder, and is skipped.
+    """
+    out = torch.zeros_like(hidden)
+    hook = getattr(_hooks, "by_module", {}).get(id(self))
+    up = "gate_up_proj" if self.has_gate else "up_proj"
+    for e in torch.unique(top_k_index).tolist():
+        if e >= self.num_experts:
+            continue
+        tokens, slot = torch.where(top_k_index == e)
+        pre = _linear(
+            hidden[tokens], expert_weight(self, up, e), _bias(self, up, e), self.is_transposed
+        )
+        act = self._apply_gate(pre) if self.has_gate else self.act_fn(pre)
+        weights = top_k_weights[tokens, slot]
+        if hook is not None:
+            act = hook(act, e, tokens, weights)
+        down = expert_weight(self, "down_proj", e)
+        y = _linear(act, down, _bias(self, "down_proj", e), self.is_transposed)
+        if self.has_post_expert_norm:
+            y = self.post_expert_norm(y)
+        out.index_add_(0, tokens, (y * weights[:, None]).to(out.dtype))
+    return out
+
+
+def _bias(experts: nn.Module, name: str, e: int) -> torch.Tensor | None:
+    return expert_weight(experts, f"{name}_bias", e) if experts.has_bias else None
+
+
+def expert_weight(experts: nn.Module, name: str, e: int) -> torch.Tensor:
+    """Expert ``e``'s slice of the stacked tensor ``name``, read from disk if streamed."""
+    param = getattr(experts, name)
+    if param.device.type != "meta":
+        return param[e]
+    weights = experts._hf_hook.weights_map  # streamed: the tensor stays on disk
+    return weights.dataset.slice(weights.prefix + name, e).to(param.dtype)
+
+
+@contextmanager
+def expert_hook(experts: nn.Module, fn: ExpertHook) -> Iterator[None]:
+    """Pass every expert's hidden activation in ``experts`` through ``fn``, in this thread."""
+    by_module = getattr(_hooks, "by_module", None)
+    if by_module is None:
+        by_module = _hooks.by_module = {}
+    previous = by_module.get(id(experts))
+    by_module[id(experts)] = fn
+    try:
+        yield
+    finally:
+        if previous is None:
+            del by_module[id(experts)]
+        else:
+            by_module[id(experts)] = previous
+
+
+def experts_modules(model: nn.Module) -> list[nn.Module]:
+    """The experts modules that ``use_experts_implementation`` decorated, in model order."""
+    return [m for m in model.modules() if hasattr(m, "has_gate") and hasattr(m, "down_proj")]
+
+
+def use_openinterp_experts(model: nn.Module) -> bool:
+    """Run ``model``'s experts with :func:`experts_forward`; False if it has none.
+
+    In a streamed model each experts module's tensors then stay on disk between passes.
+    Each pass reads only the routed experts' slices of them.
+    """
+    experts = experts_modules(model)
+    if not experts:
+        return False
+    _register()
+    model.set_experts_implementation(EXPERTS_IMPLEMENTATION)
+    for module in experts:
+        hook = getattr(module, "_hf_hook", None)
+        if hook is not None and hook.offload:
+            hook.offload = False  # the forward reads slices: don't load the whole stack
+    return True
+
+
+def experts_implementation(model: nn.Module) -> str | None:
+    """The experts forward ``model`` runs with, or ``None`` if it has no experts."""
+    if not experts_modules(model):
+        return None
+    return getattr(model.config, "_experts_implementation", None)
+
+
+def _register() -> None:
+    from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS, ExpertsInterface
+
+    if EXPERTS_IMPLEMENTATION not in ALL_EXPERTS_FUNCTIONS:
+        ExpertsInterface.register(EXPERTS_IMPLEMENTATION, experts_forward)

openinterp/harness/patcher.py

@@ -0,0 +1,322 @@
+"""Activation patching on single neuron poles.
+
+Two directions are supported:
+
+* **Denoising** (clean -> corrupted): run the corrupted prompt but inject the
+  neuron's activation cached from the clean run. The normalised recovery of the
+  clean behaviour measures *sufficiency*.
+* **Noising** (corrupted -> clean): run the clean prompt but inject the activation
+  from the corrupted run (a resample ablation). The normalised movement towards
+  the corrupted behaviour measures *necessity*.
+
+Only the requested pole is swapped. The opposite pole keeps its own value. With
+``freeze_downstream=True``, the outputs of every later MLP are pinned to their values
+from the unpatched run. This blocks MLP-mediated self-repair (the "hydra effect"), so
+the measured effect is closer to the neuron's direct effect.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, field
+from typing import Literal, Protocol
+
+import torch
+
+from openinterp.core.exceptions import ShapeMismatchError
+from openinterp.core.schemas import Pole
+from openinterp.harness.hooks import (
+    Inputs,
+    Intervention,
+    NeuronHookRunner,
+    _recorders,
+    replace_pole,
+)
+from openinterp.harness.memory import stream_rows
+
+Metric = Callable[[torch.Tensor], torch.Tensor]
+# (patched logits [1, seq, vocab], unpatched target logits [1, seq, vocab]) -> effect
+Effect = Callable[[torch.Tensor, torch.Tensor], float]
+Positions = Literal["all", "last"]
+
+
+class NeuronRef(Protocol):
+    """Anything that names one neuron pole."""
+
+    layer: int
+    neuron_idx: int
+    pole: Pole
+
+
+def token_logit_metric(token_id: int, position: int = -1) -> Metric:
+    """Logit of ``token_id`` at ``position`` (averaged over the batch)."""
+    return lambda logits: logits[:, position, token_id].mean()
+
+
+def logit_diff_metric(correct_id: int, incorrect_id: int, position: int = -1) -> Metric:
+    """Logit of ``correct_id`` minus logit of ``incorrect_id`` at ``position`` (batch mean)."""
+    return lambda logits: (
+        logits[:, position, correct_id] - logits[:, position, incorrect_id]
+    ).mean()
+
+
+def logit_direction_metric(weights: torch.Tensor, position: int = -1) -> Metric:
+    """Projection of the logits at ``position`` onto a vocabulary direction (batch mean).
+
+    ``weights`` is centred first, so the shared logit offset (which softmax ignores)
+    cannot move the metric. With weights 1/k on k tokens this is the mean logit of
+    those tokens minus the mean logit over the vocabulary.
+    """
+    w = weights.float()
+    w = w - w.mean()
+    return lambda logits: (logits[:, position].float() @ w.to(logits.device)).mean()
+
+
+def metric_change(metric: Metric) -> Effect:
+    """Signed change in ``metric`` caused by the patch."""
+    return lambda patched, target: float(metric(patched)) - float(metric(target))
+
+
+def next_token_kl(position: int = -1) -> Effect:
+    """KL(unpatched ‖ patched) between next-token distributions at ``position``."""
+
+    def effect(patched: torch.Tensor, target: torch.Tensor) -> float:
+        p = torch.log_softmax(target[:, position].float(), dim=-1)
+        q = torch.log_softmax(patched[:, position].float(), dim=-1)
+        return max(0.0, float((p.exp() * (p - q)).sum(-1).mean()))
+
+    return effect
+
+
+def _constant(value: torch.Tensor) -> Intervention:
+    return lambda _out: value
+
+
+def _normalised(patched: float, baseline: float, reference: float, eps: float = 1e-6) -> float:
+    denom = reference - baseline
+    return math.nan if abs(denom) < eps else (patched - baseline) / denom
+
+
+@dataclass(frozen=True)
+class PatchResult:
+    """Outcome of patching a neuron pole from a source run into a target run.
+
+    ``baseline`` is the target run's metric, ``reference`` the source run's, and
+    ``patched`` the target run with the source activation injected.
+    """
+
+    baseline_metric: float
+    reference_metric: float
+    patched_metric: float
+    logit_delta: torch.Tensor  # patched - baseline logits at the final position, [vocab]
+
+    @property
+    def metric_delta(self) -> float:
+        """The patched metric minus the baseline metric."""
+        return self.patched_metric - self.baseline_metric
+
+    @property
+    def recovery(self) -> float:
+        """Fraction of the source-vs-target metric gap closed by the patch (NaN if no gap)."""
+        return _normalised(self.patched_metric, self.baseline_metric, self.reference_metric)
+
+
+@dataclass(frozen=True)
+class AblationResult:
+    """The patch results of one ablation over several source prompts."""
+
+    per_source: list[PatchResult] = field(default_factory=list)
+
+    @property
+    def mean_metric_delta(self) -> float:
+        """The mean metric delta over the source prompts."""
+        return sum(r.metric_delta for r in self.per_source) / len(self.per_source)
+
+    @property
+    def mean_recovery(self) -> float:
+        """The mean recovery over the prompts where it is defined; NaN when none is."""
+        finite = [r.recovery for r in self.per_source if math.isfinite(r.recovery)]
+        return sum(finite) / len(finite) if finite else math.nan
+
+
+class ActivationPatcher:
+    """Patch one neuron pole's activation between two runs (see the module docstring)."""
+
+    def __init__(self, runner: NeuronHookRunner) -> None:
+        self.runner = runner
+
+    def patch(
+        self,
+        source: Inputs,
+        target: Inputs,
+        neuron: NeuronRef,
+        *,
+        metric: Metric | None = None,
+        positions: Positions = "all",
+        freeze_downstream: bool = False,
+    ) -> PatchResult:
+        """Inject ``neuron``'s pole activation from the ``source`` run into the ``target`` run.
+
+        If ``metric`` is omitted it defaults to the logit of the source run's
+        top-1 next token at the final position.
+        """
+        runner = self.runner
+        layer, idx, pole = neuron.layer, neuron.neuron_idx, neuron.pole
+        downstream = list(range(layer + 1, runner.n_layers)) if freeze_downstream else []
+
+        source_logits, source_cache = runner.run_with_cache(source, [layer])
+        target_enc = runner.encode(target)
+        target_logits, target_cache, downstream_out = self._target_pass(
+            target_enc, downstream, [layer]
+        )
+
+        src_act, tgt_act = source_cache[layer], target_cache[layer]
+        if positions == "all" and src_act.shape[:2] != tgt_act.shape[:2]:
+            raise ShapeMismatchError(
+                f"source {tuple(src_act.shape[:2])} and target {tuple(tgt_act.shape[:2])} "
+                "differ in batch/sequence length; use positions='last' or align the prompts."
+            )
+
+        def inject(act: torch.Tensor) -> torch.Tensor:
+            act = act.clone()
+            if positions == "all":
+                act[..., idx] = replace_pole(act[..., idx], src_act[..., idx], pole)
+            else:
+                act[:, -1, idx] = replace_pole(act[:, -1, idx], src_act[:, -1, idx], pole)
+            return act
+
+        frozen: dict[int, Intervention] = {
+            i: (lambda _out, cached=cached: cached) for i, cached in downstream_out.items()
+        }
+        patched_logits = runner.run_with_interventions(target_enc, {layer: inject}, frozen)
+
+        if metric is None:
+            metric = token_logit_metric(int(source_logits[0, -1].argmax()))
+        return PatchResult(
+            baseline_metric=float(metric(target_logits)),
+            reference_metric=float(metric(source_logits)),
+            patched_metric=float(metric(patched_logits)),
+            logit_delta=(patched_logits[0, -1] - target_logits[0, -1]).detach().cpu(),
+        )
+
+    def patch_many(
+        self,
+        source: Inputs,
+        target: Inputs,
+        layer: int,
+        neuron_indices: Sequence[int],
+        pole: Pole,
+        *,
+        effect: Effect,
+        positions: Positions = "all",
+        freeze_downstream: bool = False,
+        batch_size: int = 8,
+    ) -> list[float]:
+        """``effect`` of patching each neuron in ``layer``, one at a time, from source to target.
+
+        With ``effect=metric_change(metric)`` this is :meth:`patch`'s ``metric_delta``.
+
+        Each neuron gets its own row of a batch, so ``k`` neurons cost ``k / batch_size``
+        forward passes instead of ``k``. A streamed model runs more rows per pass
+        (:func:`~openinterp.harness.memory.stream_rows`). ``source`` and ``target`` must be
+        single prompts.
+        """
+        runner = self.runner
+        downstream = list(range(layer + 1, runner.n_layers)) if freeze_downstream else []
+        src_act = runner.layer_activations(source, [layer])[layer]  # stops at the layer
+        target_enc = runner.encode(target)
+        target_logits, target_cache, downstream_out = self._target_pass(
+            target_enc, downstream, [layer]
+        )
+        batch_size = stream_rows(runner, target_enc["input_ids"].shape[1], batch_size, logits=True)
+        if src_act.shape[0] != 1 or target_cache[layer].shape[0] != 1:
+            raise ValueError("patch_many takes one source and one target prompt")
+        if positions == "all" and src_act.shape[:2] != target_cache[layer].shape[:2]:
+            raise ShapeMismatchError(
+                f"source {tuple(src_act.shape[:2])} and target "
+                f"{tuple(target_cache[layer].shape[:2])} differ in sequence length; "
+                "use positions='last' or align the prompts."
+            )
+        span = slice(None) if positions == "all" else slice(-1, None)
+
+        deltas: list[float] = []
+        for start in range(0, len(neuron_indices), batch_size):
+            chunk = list(neuron_indices[start : start + batch_size])
+            batch = {k: v.repeat(len(chunk), *[1] * (v.dim() - 1)) for k, v in target_enc.items()}
+
+            def inject(act: torch.Tensor, chunk=chunk) -> torch.Tensor:
+                act = act.clone()
+                for row, idx in enumerate(chunk):
+                    act[row, span, idx] = replace_pole(
+                        act[row, span, idx], src_act[0, span, idx], pole
+                    )
+                return act
+
+            n = len(chunk)
+            frozen = {i: _constant(cached.repeat(n, 1, 1)) for i, cached in downstream_out.items()}
+            patched = runner.run_with_interventions(batch, {layer: inject}, frozen)
+            deltas += [effect(patched[row : row + 1], target_logits) for row in range(len(chunk))]
+        return deltas
+
+    def _target_pass(
+        self, encoded, downstream: Sequence[int], layers: Sequence[int] = ()
+    ) -> tuple[torch.Tensor, dict[int, torch.Tensor], dict[int, torch.Tensor]]:
+        """One pass over the target: its logits, ``layers``' activations, ``downstream``'s outputs.
+
+        Patching freezes the downstream outputs at their target values. They were read in a
+        second pass over the target. One pass gives the same values for one read of the
+        layers, which is what a streamed model pays for.
+        """
+        acts: dict[int, torch.Tensor] = {}
+        outputs: dict[int, torch.Tensor] = {}
+
+        def recorder(layer: int) -> Intervention:
+            def record(out: torch.Tensor) -> None:
+                outputs[layer] = out.detach().clone()
+
+            return record
+
+        logits = self.runner.run_with_interventions(
+            encoded, _recorders(acts, layers), {i: recorder(i) for i in downstream}
+        )
+        return logits, acts, outputs
+
+    def denoise(self, clean: Inputs, corrupted: Inputs, neuron: NeuronRef, **kw) -> PatchResult:
+        """Clean -> corrupted patch. ``recovery`` is the neuron's sufficiency."""
+        return self.patch(clean, corrupted, neuron, **kw)
+
+    def noise(self, clean: Inputs, corrupted: Inputs, neuron: NeuronRef, **kw) -> PatchResult:
+        """Corrupted -> clean patch. ``recovery`` is the neuron's necessity."""
+        return self.patch(corrupted, clean, neuron, **kw)
+
+    def resample_ablation(
+        self,
+        target: Inputs,
+        resample_sources: Sequence[Inputs],
+        neuron: NeuronRef,
+        *,
+        metric: Metric | None = None,
+        positions: Positions = "last",
+        freeze_downstream: bool = False,
+    ) -> AblationResult:
+        """Replace the neuron pole in ``target`` with its value on each resample source.
+
+        The default metric is fixed from the *target* run, so every source is scored
+        against the same token.
+        """
+        if metric is None:
+            target_logits = self.runner.run_with_interventions(target)
+            metric = token_logit_metric(int(target_logits[0, -1].argmax()))
+        results = [
+            self.patch(
+                src,
+                target,
+                neuron,
+                metric=metric,
+                positions=positions,
+                freeze_downstream=freeze_downstream,
+            )
+            for src in resample_sources
+        ]
+        return AblationResult(per_source=results)

openinterp/harness/patchscope.py

@@ -0,0 +1,177 @@
+"""Verbalise a unit's output direction with a Patchscope (Ghandeharioun et al., 2024).
+
+Take the residual stream after layer ℓ's block, at a placeholder token of an inspection
+prompt. It is replaced by the unit's output direction, scaled to a multiple of the norm
+that it replaces. The model then continues greedily, and so describes the direction without
+seeing a single example. A random direction injected the same way is the control: any
+text the two share comes from the prompt, not the unit.
+
+There are two inspection prompts: a few-shot "word: category" list and a definition
+frame. Each is used at the unit's own layer ℓ (source layer = target layer). Each is also
+used at an early layer, as SelfIE (Chen et al., 2024) does. On Qwen2.5-0.5B, the few-shot
+list injected at layer 2 names the concept of several VERIFIED poles. It gives "crop",
+"newspaper", "official" and "culture", where the control gives "animal", "emotion" or
+"period".
+Injected at their own layer, poles from layer 17 up mostly read back the control's
+text. The decoding is greedy, so a reading is deterministic given the weights.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+
+import torch
+
+from openinterp.core.exceptions import ArchitectureNotSupportedError
+from openinterp.core.patchscope import PatchscopeReading, PatchscopeReadout
+from openinterp.harness.hooks import NeuronHookRunner
+from openinterp.harness.units import FeatureUnit, NeuronUnit, Unit
+
+PROMPTS = {
+    "category": "banana: fruit\nhammer: tool\nviolin: music\nsenator: politics\n{x}:",
+    "meaning": 'The meaning of the word "{x}" is',
+}
+PLACEHOLDER = "?"
+SCALES = (1.0, 2.0, 4.0)
+EARLY_LAYER = 2
+MAX_NEW_TOKENS = 12
+CONTROL_SEED = 0
+
+
+def write_direction(runner: NeuronHookRunner, unit: Unit) -> torch.Tensor | None:
+    """Return the vector that one unit of the unit's strength adds to the residual stream.
+
+    That is a pole's signed ``w_out``, or a dictionary latent's decoder row. A head writes
+    a subspace, not one direction, so it has none.
+    """
+    if isinstance(unit, NeuronUnit):
+        ident = unit.ident
+        direction = runner.neuron_output_direction(ident.layer, ident.neuron_idx).float()
+        return -direction if ident.pole == "negative" else direction
+    if isinstance(unit, FeatureUnit):
+        return unit.dictionary(runner).w_dec[unit.ident.feature_idx].float()
+    return None
+
+
+def _prompt_ids(runner: NeuronHookRunner, template: str, placeholder: str) -> tuple[list[int], int]:
+    """Token ids of ``template`` with the placeholder as exactly one token, and its index."""
+    tok = runner.tokenizer
+    before, after = template.split("{x}")
+    marker = tok.encode(placeholder, add_special_tokens=False)
+    if len(marker) != 1:
+        raise ValueError(f"the placeholder {placeholder!r} must be one token, not {len(marker)}")
+    prefix = tok.encode(before, add_special_tokens=True) if before else []
+    return prefix + marker + tok.encode(after, add_special_tokens=False), len(prefix)
+
+
+def _block_output(output) -> torch.Tensor:
+    return output[0] if isinstance(output, tuple) else output
+
+
+@torch.no_grad()
+def _residual_norm(runner: NeuronHookRunner, ids: list[int], layer: int, pos: int) -> float:
+    captured: list[float] = []
+
+    def capture(_module, _args, output):
+        captured.append(float(_block_output(output)[0, pos].float().norm()))
+
+    handle = runner.layers[layer].register_forward_hook(capture)
+    try:
+        runner.forward({"input_ids": torch.tensor([ids], device=runner.device)})
+    finally:
+        handle.remove()
+    if not captured:
+        raise ArchitectureNotSupportedError(f"layer {layer}'s block never ran as a module")
+    return captured[-1]
+
+
+@torch.no_grad()
+def _continue(
+    runner: NeuronHookRunner,
+    ids: list[int],
+    pos: int,
+    layer: int,
+    vectors: torch.Tensor,
+    max_new_tokens: int,
+) -> list[str]:
+    """Return greedy continuations of ``ids``, one per row of ``vectors``.
+
+    In each, the residual after ``layer``'s block at ``pos`` is replaced by that row. Each
+    stops at a line break or the end-of-text token.
+    """
+    tok = runner.tokenizer
+    eos = getattr(tok, "eos_token_id", None)
+
+    def patch(_module, _args, output):
+        hidden = _block_output(output).clone()
+        hidden[:, pos] = vectors.to(hidden.device, hidden.dtype)
+        return (hidden, *output[1:]) if isinstance(output, tuple) else hidden
+
+    seq = torch.tensor([ids] * len(vectors), device=runner.device)
+    done = [False] * len(vectors)
+    handle = runner.layers[layer].register_forward_hook(patch)
+    try:
+        for _ in range(max_new_tokens):
+            nxt = runner.forward({"input_ids": seq})[:, -1].argmax(-1)
+            seq = torch.cat([seq, nxt[:, None]], dim=1)
+            for row, token in enumerate(nxt.tolist()):
+                done[row] |= token == eos or "\n" in tok.decode([token])
+            if all(done):
+                break
+    finally:
+        handle.remove()
+    texts = []
+    for row in seq[:, len(ids) :].tolist():
+        if eos in row:
+            row = row[: row.index(eos)]
+        texts.append(tok.decode(row).split("\n")[0].strip())
+    return texts
+
+
+def patchscope(
+    runner: NeuronHookRunner,
+    layer: int,
+    direction: torch.Tensor,
+    *,
+    prompts: Mapping[str, str] = PROMPTS,
+    placeholder: str = PLACEHOLDER,
+    scales: Sequence[float] = SCALES,
+    early_layer: int = EARLY_LAYER,
+    max_new_tokens: int = MAX_NEW_TOKENS,
+    control_seed: int = CONTROL_SEED,
+) -> PatchscopeReadout:
+    """Read ``direction``, written after ``layer``'s MLP, through every inspection prompt.
+
+    It is read at ``layer``, and at ``early_layer`` when that is earlier, at each scale.
+    """
+    if runner.tokenizer is None:
+        raise ValueError("a Patchscope needs a tokenizer")
+    direction = direction.float().cpu()
+    generator = torch.Generator().manual_seed(control_seed)
+    control = torch.randn(direction.shape, generator=generator)
+    units = torch.stack([direction / direction.norm(), control / control.norm()])
+    readings = []
+    for name, template in prompts.items():
+        ids, pos = _prompt_ids(runner, template, placeholder)
+        for target in sorted({min(early_layer, layer), layer}):
+            norm = _residual_norm(runner, ids, target, pos)
+            # Rows: the unit at each scale, then the control at each scale.
+            vectors = torch.stack([units[i] * norm * s for i in (0, 1) for s in scales])
+            texts = _continue(runner, ids, pos, target, vectors, max_new_tokens)
+            for i, scale in enumerate(scales):
+                readings.append(
+                    PatchscopeReading(
+                        prompt=name,
+                        layer=target,
+                        scale=float(scale),
+                        text=texts[i],
+                        control_text=texts[len(scales) + i],
+                    )
+                )
+    return PatchscopeReadout(
+        prompts=dict(prompts),
+        placeholder=placeholder,
+        max_new_tokens=max_new_tokens,
+        control_seed=control_seed,
+        readings=readings,
+    )