feat(steer): name the steering methods once, as SteerMethod

docs/INTERNALS.md

@@ -54,7 +54,9 @@ files and why they are where they are.
 - `steer.py` — steering, and the per-token generation stream. Each method's arithmetic is one
   `steer_delta` branch, which is what the vLLM worker's modifier computes too and what
   `tests/test_steer_math_parity.py` runs against it on CPU; `steer()` is the context both backends
-  take, registering per-request on vLLM rather than installing a global hook.
+  take, registering per-request on vLLM rather than installing a global hook. The methods are named
+  once, in `steer_specs.SteerMethod`: the eager `SteerSpec.method` and the worker dict's `op` both
+  read from it, so a method has one spelling on either backend.
 - `mappers.py` — translation between canonical points and other frameworks' names:
   TransformerLens hook strings and nnsight/nnterp accessors, both directions. See [Porting from
   TransformerLens, nnsight or nnterp](PORTING.md#porting-from-transformerlens-nnsight-or-nnterp).

interp_engine/__init__.py

@@ -120,6 +120,8 @@
     ProjectionCapSpec,
     SteeringOp,
     SteeringSpec,
+    SteerMethod,
+    steer_method,
     steering_spec_to_worker_specs,
 )
 from interp_engine.sync import SyncModel, sync_model
@@ -245,13 +247,15 @@
     "PositionMask",
     "ProjectionCapSpec",
     "SteerMask",
+    "SteerMethod",
     "SteerSpec",
     "SteeringOp",
     "SteeringSpec",
     "projection_cap_delta",
     "resolve_masked_positions",
     "steer",
     "steer_delta",
+    "steer_method",
     "steering_spec_to_eager_specs",
     "steering_spec_to_worker_specs",
     "unit_vector",

interp_engine/steer.py

@@ -32,6 +32,8 @@
     OrthogonalDecompSpec,
     ProjectionCapSpec,
     SteeringSpec,
+    SteerMethod,
+    steer_method,
 )
 from interp_engine.sync import sync_model
 
@@ -156,10 +158,10 @@ def projection_cap_delta(
     return (capped - projection) * unit
 
 
-#: The steering methods :class:`SteerSpec` accepts, each with the same arithmetic as the
-#: worker-side op of the same name (``vllm_capture/steering.py``). Kept as a tuple so the
-#: refusal for an unknown method can list them.
-STEER_METHODS = ("additive", "orthogonal", "projection_cap")
+#: The steering methods :class:`SteerSpec` accepts, as their spellings. Derived from
+#: :class:`~interp_engine.steer_specs.SteerMethod`, which the worker-side op reads from too, and
+#: kept under this name so callers that listed the methods keep working.
+STEER_METHODS: tuple[str, ...] = tuple(m.value for m in SteerMethod)
 
 
 @dataclass
@@ -169,7 +171,7 @@ class SteerSpec:
     vector: torch.Tensor
     layer: int
     coeff: float = 1.0
-    method: str = "additive"  # one of STEER_METHODS
+    method: SteerMethod = SteerMethod.ADDITIVE
     point: str = "resid_post"
     normalize: bool = False
     stream: int | None = None
@@ -195,6 +197,11 @@ class SteerSpec:
     max: float | None = None
     """Upper bound for ``method="projection_cap"``. See :attr:`min`."""
 
+    def __post_init__(self) -> None:
+        # A spec built from a string -- a config file, a wire payload -- is checked here, at
+        # construction, rather than at the first forward it reaches. An unknown name lists the members.
+        self.method = steer_method(self.method)
+
 
 def _prepared_vector(spec: SteerSpec, ref: torch.Tensor) -> torch.Tensor:
     vec = spec.vector.to(device=ref.device, dtype=ref.dtype)
@@ -214,13 +221,13 @@ def steer_delta(spec: SteerSpec, activations: torch.Tensor, vector: torch.Tensor
     orthogonal case subtly wrong if either half is forgotten. It is also the shape the vLLM
     worker's modifiers already have, so the two backends run the same expressions.
     """
-    if spec.method == "additive":
-        return spec.coeff * vector
-    if spec.method == "orthogonal":
-        return OrthogonalProjector(vector).delta(activations, spec.coeff)
-    if spec.method == "projection_cap":
-        return projection_cap_delta(activations, vector, minimum=spec.min, maximum=spec.max)
-    raise ValueError(f"Unknown steering method {spec.method!r}; expected one of {', '.join(STEER_METHODS)}")
+    match steer_method(spec.method):
+        case SteerMethod.ADDITIVE:
+            return spec.coeff * vector
+        case SteerMethod.ORTHOGONAL:
+            return OrthogonalProjector(vector).delta(activations, spec.coeff)
+        case SteerMethod.PROJECTION_CAP:
+            return projection_cap_delta(activations, vector, minimum=spec.min, maximum=spec.max)
 
 
 @dataclass(frozen=True)
@@ -420,20 +427,24 @@ def steering_spec_to_eager_specs(spec: SteeringSpec, *, point: str | None = None
             if isinstance(op, AddSpec):
                 vector = op.vector if isinstance(op.vector, torch.Tensor) else torch.tensor(op.vector)
                 out.append(
-                    SteerSpec(vector=vector, layer=int(layer), coeff=float(op.scale), method="additive", **where)
+                    SteerSpec(
+                        vector=vector, layer=int(layer), coeff=float(op.scale), method=SteerMethod.ADDITIVE, **where
+                    )
                 )
             elif isinstance(op, OrthogonalDecompSpec):
                 vector = op.vector if isinstance(op.vector, torch.Tensor) else torch.tensor(op.vector)
                 out.append(
-                    SteerSpec(vector=vector, layer=int(layer), coeff=float(op.coeff), method="orthogonal", **where)
+                    SteerSpec(
+                        vector=vector, layer=int(layer), coeff=float(op.coeff), method=SteerMethod.ORTHOGONAL, **where
+                    )
                 )
             elif isinstance(op, ProjectionCapSpec):
                 vector = op.vector if isinstance(op.vector, torch.Tensor) else torch.tensor(op.vector)
                 out.append(
                     SteerSpec(
                         vector=vector,
                         layer=int(layer),
-                        method="projection_cap",
+                        method=SteerMethod.PROJECTION_CAP,
                         min=op.min,
                         max=op.max,
                         **where,

interp_engine/steer_specs.py

@@ -9,10 +9,43 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+from enum import StrEnum
 
 import torch
 
 
+class SteerMethod(StrEnum):
+    """The steering arithmetics, named once for both backends.
+
+    The eager ``SteerSpec.method``, the worker dict's ``op`` and the static tap's op set all read from
+    here, so a method has one spelling on every path. It used to have two -- ``additive`` on the
+    eager spec and ``add`` on the worker dict -- and nothing but a runtime set-membership check told
+    a caller which one it was speaking to. A ``StrEnum`` so a member compares equal to its own
+    spelling: a spec loaded from JSON with ``"orthogonal"`` in it needs no translation.
+
+    A new method is a new member here, plus a delta in ``steer.py`` and a modifier in
+    ``vllm_capture/steering.py`` -- see ``AGENTS.md``, "Steering arithmetic is one function per
+    method, shared with the worker".
+    """
+
+    ADDITIVE = "additive"
+    ORTHOGONAL = "orthogonal"
+    PROJECTION_CAP = "projection_cap"
+
+
+def steer_method(value: SteerMethod | str) -> SteerMethod:
+    """``value`` as a :class:`SteerMethod`, refusing anything else with the members listed.
+
+    The enum's own ``ValueError`` says only that the value is invalid; a caller who wrote ``"add"``
+    is better served by seeing ``additive`` in the message.
+    """
+    try:
+        return SteerMethod(value)
+    except ValueError:
+        expected = ", ".join(m.value for m in SteerMethod)
+        raise ValueError(f"Unknown steering method {value!r}; expected one of {expected}") from None
+
+
 def _as_list(vector: torch.Tensor | list[float]) -> list[float]:
     if isinstance(vector, torch.Tensor):
         return vector.detach().float().flatten().tolist()
@@ -102,7 +135,7 @@ def steering_spec_to_worker_specs(spec: SteeringSpec, *, point: str | None = Non
                     {
                         "layer": int(layer),
                         **where,
-                        "op": "add",
+                        "op": SteerMethod.ADDITIVE.value,
                         "vector": _as_list(op.vector),
                         "coeff": float(op.scale),
                     }
@@ -112,7 +145,7 @@ def steering_spec_to_worker_specs(spec: SteeringSpec, *, point: str | None = Non
                     {
                         "layer": int(layer),
                         **where,
-                        "op": "projection_cap",
+                        "op": SteerMethod.PROJECTION_CAP.value,
                         "vector": _as_list(op.vector),
                         "min": op.min,
                         "max": op.max,
@@ -123,7 +156,7 @@ def steering_spec_to_worker_specs(spec: SteeringSpec, *, point: str | None = Non
                     {
                         "layer": int(layer),
                         **where,
-                        "op": "orthogonal",
+                        "op": SteerMethod.ORTHOGONAL.value,
                         "vector": _as_list(op.vector),
                         "coeff": float(op.coeff),
                     }

interp_engine/vllm_backend.py

@@ -34,6 +34,7 @@
 from interp_engine.points import d_model_wide, hyper_connection_names, refusal_reasons
 from interp_engine.points import steer_refusal_reason as points_steer_refusal
 from interp_engine.residual_basis import ResidualBasis, vllm_residual_basis
+from interp_engine.steer_specs import SteerMethod
 from interp_engine.vllm_capture import (
     _GLOBAL_POINTS,
     DEFAULT_HS_STORAGE_PATH,
@@ -52,6 +53,7 @@
     DECODE_ONLY_GRAPHS,
     STATIC_ENV,
     STATIC_SKIP_ABSENT_ENV,
+    STATIC_WRITE_OPS,
     apply_breakable_env,
     decode_only_graphs_reason,
     encode_static_env,
@@ -1504,20 +1506,20 @@ def _require_capture_points(self, pts: Sequence[str], what: str) -> None:
     def _require_static_writes(self, specs: Sequence[dict], what: str) -> None:
         """Allow static writes when hooks run, or when every write site is declared.
 
-        Additive ``add_`` and live-read ops (orthogonal, projection_cap, lens steer/ablate/swap)
-        all ride the same static wrap. A site miss is a 400, not a silent no-op.
+        The additive static ``add_`` and live-read ops (orthogonal, projection_cap, lens
+        steer/ablate/swap) all ride the same static wrap. A site miss is a 400, not a silent no-op.
         """
         if self.hooks_available:
             return
         writes = getattr(self, "_static_writes", frozenset())
         if self.graph_replay and writes:
             for spec in specs:
-                op = spec.get("op", "add")
-                if op not in {"add", "orthogonal", "projection_cap", "steer", "ablate", "swap"}:
+                op = spec.get("op", SteerMethod.ADDITIVE)
+                if op not in STATIC_WRITE_OPS:
                     raise RuntimeError(
                         f"{what} op {op!r} is not one a static write tap can apply; "
-                        f"backend='vllm-static' serves add, orthogonal, projection_cap, steer, "
-                        f"ablate and swap. Use backend='vllm' for this one."
+                        f"backend='vllm-static' serves {', '.join(sorted(STATIC_WRITE_OPS))}. "
+                        f"Use backend='vllm' for this one."
                     )
                 site = Address(str(spec.get("point") or "resid_post"), int(spec["layer"]))
                 if not any(alias in writes for alias in resid_stream_aliases(site)):

interp_engine/vllm_capture/static.py

@@ -25,6 +25,7 @@
 from interp_engine.facts import is_linear_attention_layer, unclassified_layer_kinds
 from interp_engine.hooks import hidden_arg_index, hidden_from_call
 from interp_engine.points import steer_refusal_reason
+from interp_engine.steer_specs import SteerMethod
 from interp_engine.vllm_capture._demux import _ensure_patched, _get_demux, _resolve_rid
 from interp_engine.vllm_capture._hooks import (
     flat_value,
@@ -675,10 +676,10 @@ def fit_max_num_batched_tokens(
 # --- worker state ------------------------------------------------------------
 
 
-#: Ops static writes can serve. ``add`` uses a static ``add_``; the rest read the live residual.
-_STEER_OPS = frozenset({"add", "orthogonal", "projection_cap"})
+#: Ops static writes can serve. ``additive`` uses a static ``add_``; the rest read the live residual.
+_STEER_OPS = frozenset(m.value for m in SteerMethod)
 _LENS_OPS = frozenset({"steer", "ablate", "swap"})
-_STATIC_WRITE_OPS = _STEER_OPS | _LENS_OPS
+STATIC_WRITE_OPS = _STEER_OPS | _LENS_OPS
 
 
 @dataclass
@@ -1811,7 +1812,7 @@ def worker_set_static_delta(
 ) -> None:
     """Install static writes from worker specs. Zeros every write site first.
 
-    Additive ``op="add"`` without a lens scope fills the static ``delta`` buffer. Orthogonal,
+    ``op="additive"`` without a lens scope fills the static ``delta`` buffer. Orthogonal,
     projection_cap, and lens ops attach a live ``modify`` that reads the residual each
     forward (breakable ``add_eager``). ``lens_scope`` is the jlens prefill/decode skip.
 
@@ -1830,12 +1831,12 @@ def worker_set_static_delta(
     from interp_engine.vllm_capture.steering import _make_steer_modifier
 
     for spec in specs:
-        op = str(spec.get("op", "add"))
-        if op not in _STATIC_WRITE_OPS:
-            raise ValueError(f"static taps cannot serve op={op!r}; supported ops are {sorted(_STATIC_WRITE_OPS)}")
+        op = str(spec.get("op", SteerMethod.ADDITIVE))
+        if op not in STATIC_WRITE_OPS:
+            raise ValueError(f"static taps cannot serve op={op!r}; supported ops are {sorted(STATIC_WRITE_OPS)}")
         site = _write_site(static, Address(str(spec["point"]), int(spec["layer"])))
         assert site.delta is not None
-        if op == "add" and lens_scope is None and spec.get("stream") is None:
+        if op == SteerMethod.ADDITIVE and lens_scope is None and spec.get("stream") is None:
             vec = torch.tensor(spec["vector"], dtype=torch.float32, device=site.delta.device)
             vec = (vec * float(spec["coeff"])).to(dtype=site.delta.dtype)
             site.delta.copy_(vec.reshape(1, -1).expand_as(site.delta))
@@ -1857,14 +1858,14 @@ def _compile_write_req(
     prompt_len: int,
     steer_generated: bool,
 ) -> _WriteReq:
-    op = str(spec.get("op", "add"))
-    if op not in _STATIC_WRITE_OPS:
-        raise ValueError(f"static taps cannot serve op={op!r}; supported ops are {sorted(_STATIC_WRITE_OPS)}")
+    op = str(spec.get("op", SteerMethod.ADDITIVE))
+    if op not in STATIC_WRITE_OPS:
+        raise ValueError(f"static taps cannot serve op={op!r}; supported ops are {sorted(STATIC_WRITE_OPS)}")
     assert site.delta is not None
     device, dtype = site.delta.device, site.delta.dtype
     # A constant `[1, width]` vector broadcasts over a stream axis and so cannot exclude one; see
     # `worker_set_static_delta` for why a `stream` therefore has to go the modifier way.
-    if op == "add" and spec.get("stream") is None:
+    if op == SteerMethod.ADDITIVE and spec.get("stream") is None:
         vec = torch.tensor(spec["vector"], dtype=torch.float32, device=device)
         vec = (vec * float(spec["coeff"])).to(dtype=dtype).reshape(1, -1)
         return _WriteReq(