RoutedDecoding

steerability.algorithms.output_control.routed_decoding

actions

Route actions for RoutedDecoding, each lowered to a PhasedDriver phase plan.

Generate dataclass

Plain pass-through: delegate the row to the model with no splicing.

Source code in steerability/algorithms/output_control/routed_decoding/actions.py

@dataclass(frozen=True)
class Generate:
    """Plain pass-through: delegate the row to the model with no splicing."""

    def plan(self) -> list:
        """The phase plan: one unbounded `Generated` phase."""
        return [Generated()]

    def __str__(self) -> str:
        return "generate"

plan()

The phase plan: one unbounded Generated phase.

Source code in steerability/algorithms/output_control/routed_decoding/actions.py

def plan(self) -> list:
    """The phase plan: one unbounded `Generated` phase."""
    return [Generated()]

Prefix dataclass

Canned prefix followed by generation: splice text, then delegate to the model.

Attributes:

Name	Type	Description
text	str	The prefix text, tokenized without special tokens.

Source code in steerability/algorithms/output_control/routed_decoding/actions.py

@dataclass(frozen=True)
class Prefix:
    """Canned prefix followed by generation: splice `text`, then delegate to the model.

    Attributes:
        text: The prefix text, tokenized without special tokens.
    """

    text: str

    def __post_init__(self):
        if not isinstance(self.text, str) or not self.text:
            raise ValueError("Prefix requires a non-empty text.")

    def plan(self) -> list:
        """The phase plan: an appended `Fixed` phase, then an unbounded `Generated` phase."""
        return [Fixed(self.text, add_special_tokens=False), Generated()]

    def __str__(self) -> str:
        return f'prefix("{_ellipsize(self.text)}") + generate'

text instance-attribute

plan()

The phase plan: an appended Fixed phase, then an unbounded Generated phase.

Source code in steerability/algorithms/output_control/routed_decoding/actions.py

def plan(self) -> list:
    """The phase plan: an appended `Fixed` phase, then an unbounded `Generated` phase."""
    return [Fixed(self.text, add_special_tokens=False), Generated()]

Respond dataclass

Terminal canned response: splice text into the stream and generate nothing.

Attributes:

Name	Type	Description
text	str	The canned response text, tokenized without special tokens.

Source code in steerability/algorithms/output_control/routed_decoding/actions.py

@dataclass(frozen=True)
class Respond:
    """Terminal canned response: splice `text` into the stream and generate nothing.

    Attributes:
        text: The canned response text, tokenized without special tokens.
    """

    text: str

    def __post_init__(self):
        if not isinstance(self.text, str) or not self.text:
            raise ValueError("Respond requires a non-empty text.")

    def plan(self) -> list:
        """The phase plan: one appended `Fixed` phase."""
        return [Fixed(self.text, add_special_tokens=False)]

    def __str__(self) -> str:
        return f'respond("{_ellipsize(self.text)}")'

text instance-attribute

plan()

The phase plan: one appended Fixed phase.

Source code in steerability/algorithms/output_control/routed_decoding/actions.py

def plan(self) -> list:
    """The phase plan: one appended `Fixed` phase."""
    return [Fixed(self.text, add_special_tokens=False)]

generate()

A plain pass-through action.

Source code in steerability/algorithms/output_control/routed_decoding/actions.py

def generate() -> Generate:
    """A plain pass-through action."""
    return Generate()

prefix(text)

A canned-prefix-then-generate action carrying text.

Source code in steerability/algorithms/output_control/routed_decoding/actions.py

def prefix(text: str) -> Prefix:
    """A canned-prefix-then-generate action carrying `text`."""
    return Prefix(text)

respond(text)

A terminal canned-response action carrying text.

Source code in steerability/algorithms/output_control/routed_decoding/actions.py

def respond(text: str) -> Respond:
    """A terminal canned-response action carrying `text`."""
    return Respond(text)

args

RoutedDecoding argument validation.

RoutedDecodingArgs dataclass

Bases: BaseArgs

Arguments for RoutedDecoding.

Attributes:

Name	Type	Description
probes	ProbeSet | ProbeSetFit | None	The probes whose decisions drive routing. Either a fitted ProbeSet, or a ProbeSetFit recipe the driver fits at steer() time on the model the pipeline provides.
rules	Router | None	The Router mapping probe decisions to actions. Probe names are validated here against probes.names (available in both forms), so misconfigured routes fail at construction, before any model loads.
allow_model_mismatch	bool	When True, a fitted ProbeSet whose recorded model fingerprints differ from the pipeline's model is accepted at steer() time instead of raising.

Source code in steerability/algorithms/output_control/routed_decoding/args.py

@dataclass
class RoutedDecodingArgs(BaseArgs):
    """Arguments for `RoutedDecoding`.

    Attributes:
        probes: The probes whose decisions drive routing. Either a fitted `ProbeSet`, or a
            `ProbeSetFit` recipe the driver fits at `steer()` time on the model the pipeline
            provides.
        rules: The `Router` mapping probe decisions to actions. Probe names are validated
            here against `probes.names` (available in both forms), so misconfigured routes fail
            at construction, before any model loads.
        allow_model_mismatch: When True, a fitted `ProbeSet` whose recorded model fingerprints
            differ from the pipeline's model is accepted at `steer()` time instead of raising.
    """

    probes: ProbeSet | ProbeSetFit | None = None
    rules: Router | None = None
    allow_model_mismatch: bool = False

    def __post_init__(self):
        if not isinstance(self.probes, (ProbeSet, ProbeSetFit)):
            raise TypeError(
                f"probes must be a ProbeSet or ProbeSetFit instance; got "
                f"{type(self.probes).__name__}."
            )
        if not isinstance(self.rules, Router):
            raise TypeError(
                f"rules must be a Router instance; got "
                f"{type(self.rules).__name__}."
            )
        self.rules.validate_names(set(self.probes.names))

allow_model_mismatch = False class-attribute instance-attribute

probes = None class-attribute instance-attribute

rules = None class-attribute instance-attribute

validate(_init_data=None, **kwargs) classmethod

Create and validate an Args instance from dict, kwargs, or existing instance.

Parameters:

Name	Type	Description	Default
_init_data	Any | None	Existing instance, dict of args, or None. Named with underscore prefix to avoid collision with common field names like "data".	None

Returns:

Type	Description
T	Validated instance of the Args class

Source code in steerability/algorithms/core/base_args.py

@classmethod
def validate(cls: type[T], _init_data: Any | None = None, **kwargs) -> T:
    """Create and validate an Args instance from dict, kwargs, or existing instance.

    Args:
        _init_data: Existing instance, dict of args, or None. Named with underscore prefix
            to avoid collision with common field names like "data".

    Returns:
        Validated instance of the Args class
    """

    if isinstance(_init_data, cls):
        return _init_data

    if isinstance(_init_data, Mapping):
        kwargs = {**_init_data, **kwargs}

    return cls(**kwargs)

control

RoutedDecoding control: route each row to a response strategy based on probe decisions.

RoutedDecoding

Bases: PhasedDriver

Decoding driver that routes each prompt to a response strategy via probe decisions.

RoutedDecoding pairs a ProbeSet (named calibrated probes scored in one read-only forward) with a Router (ordered routes over the probe names, first match wins, evaluated per row). Decoding proceeds in three steps:

1. Probe pass: one read-only forward of the prompt, issued by ProbeSet.read(), yields each probe's per-row signed score and decision (score >= 0).
2. Routing: the decisions feed rules.route(), matching each row to its first satisfied route (or the default).
3. Execution: each row's action is lowered to a phase plan and executed by the inherited plan runner. Respond(text) splices the canned tokens with no generation; Prefix(text) splices the prefix then generates; Generate() delegates the row to model.generate. A raw list[Fixed | Generated] is also accepted as an action. The composed logits processors and stopping criteria apply in every generated phase, so the driver contract is preserved.

The probes arrive fitted (a ProbeSet) or as a deferred recipe (a ProbeSetFit) that steer() fits on the model the pipeline provides, so pipelines whose structural controls produce the final weights fit on those weights. A fitted set whose recorded model fingerprints differ from the pipeline's model raises at steer() unless allow_model_mismatch=True.

The read wraps its forward in auxiliary_pass(aligned=True), the standard marking for same-model forwards issued during decoding (same_model_forwards = True). Its capture hooks are closures registered for the duration of that single forward and are removed before decoding begins, so the re-processing of the prompt inside generated phases is never re-scored.

Costs per row: a routed canned response costs one prompt forward and zero decode steps; a pass-through row costs one extra prompt forward compared to the default driver.

When state controls carrying behavior transforms share the pipeline, their hooks are live during the probe pass. The pass is trajectory-aligned, so "all"-scope transforms, and position-scoped ones covering prompt positions, apply to it, and probe scores are measured under that steering; token_scope="after_prompt" transforms are inert (the pass contains only prompt positions). Those controls' condition scorers, gates, and position counters ignore the pass entirely, since it is auxiliary.

Padding is handled per row: pad positions are stripped (via the attention mask) before plan execution, and each returned row is the original padded prompt plus its continuation, so the pipeline's prompt-length slicing stays exact.

The most recent routing outcome is retained on latest_routes (one route name per row, "default" for unmatched rows); per-probe decisions and scores are available on the probe set's latest readings.

runtime_kwargs:

• "canned_responses": dict mapping route names to replacement text, overriding the Respond/Prefix text of matching routes for this call only. Keys that do not name a Respond/Prefix route are ignored with a warning.

Source code in steerability/algorithms/output_control/routed_decoding/control.py

class RoutedDecoding(PhasedDriver):
    """Decoding driver that routes each prompt to a response strategy via probe decisions.

    `RoutedDecoding` pairs a `ProbeSet` (named calibrated probes scored in one read-only
    forward) with a `Router` (ordered routes over the probe names, first match wins, evaluated
    per row). Decoding proceeds in three steps:

    1. **Probe pass**: one read-only forward of the prompt, issued by `ProbeSet.read()`, yields
       each probe's per-row signed score and decision (`score >= 0`).
    2. **Routing**: the decisions feed `rules.route()`, matching each row to its first
       satisfied route (or the default).
    3. **Execution**: each row's action is lowered to a phase plan and executed by the
       inherited plan runner. `Respond(text)` splices the canned tokens with no generation;
       `Prefix(text)` splices the prefix then generates; `Generate()` delegates the row to
       `model.generate`. A raw `list[Fixed | Generated]` is also accepted as an action. The
       composed logits processors and stopping criteria apply in every generated phase, so the
       driver contract is preserved.

    The probes arrive fitted (a `ProbeSet`) or as a deferred recipe (a `ProbeSetFit`) that
    `steer()` fits on the model the pipeline provides, so pipelines whose structural controls
    produce the final weights fit on those weights. A fitted set whose recorded model
    fingerprints differ from the pipeline's model raises at `steer()` unless
    `allow_model_mismatch=True`.

    The read wraps its forward in `auxiliary_pass(aligned=True)`, the standard marking for
    same-model forwards issued during decoding (`same_model_forwards = True`). Its capture
    hooks are closures registered for the duration of that single forward and are removed
    before decoding begins, so the re-processing of the prompt inside generated phases is
    never re-scored.

    Costs per row: a routed canned response costs one prompt forward and zero decode steps; a
    pass-through row costs one extra prompt forward compared to the default driver.

    When state controls carrying behavior transforms share the pipeline, their hooks are live
    during the probe pass. The pass is trajectory-aligned, so `"all"`-scope transforms, and
    position-scoped ones covering prompt positions, apply to it, and probe scores are measured
    under that steering; `token_scope="after_prompt"` transforms are inert (the pass contains
    only prompt positions). Those controls' condition scorers, gates, and position counters
    ignore the pass entirely, since it is auxiliary.

    Padding is handled per row: pad positions are stripped (via the attention mask) before plan
    execution, and each returned row is the original padded prompt plus its continuation, so
    the pipeline's prompt-length slicing stays exact.

    The most recent routing outcome is retained on `latest_routes` (one route name per row,
    `"default"` for unmatched rows); per-probe decisions and scores are available on the probe
    set's `latest` readings.

    `runtime_kwargs`:

    - `"canned_responses"`: dict mapping route names to replacement text, overriding the
      `Respond`/`Prefix` text of matching routes for this call only. Keys that do not name a
      `Respond`/`Prefix` route are ignored with a warning.
    """

    Args = RoutedDecodingArgs

    supports_batching: bool = True
    same_model_forwards: bool = True

    RUNTIME_KWARGS_SCHEMA = [
        {
            "name": "canned_responses",
            "type": "dict[str, str]",
            "scope": "call",
            "description": "Per-call override of Respond/Prefix text, keyed by route name.",
        },
    ]

    tokenizer: PreTrainedTokenizerBase | None = None

    def __init__(self, *args, **kwargs):
        # route through OutputControl (validate RoutedDecodingArgs, mirror fields, _configure)
        OutputControl.__init__(self, *args, **kwargs)

    def _configure(self) -> None:
        """Initialize the `PhasedDriver` fields the plan runner reads."""
        self.extract_after = None
        self.tokenizer = None
        self.latest_routes: list[str] = []

    def max_rollouts_per_query(self) -> int:
        """1: every route lowers to at most one `Generated` phase, and the probe pass is a
        read-only capture rather than a generation."""
        return 1

    def requirements(self) -> Requirements:
        """In-process torch or hidden-state capture at generate; the probe pass reads the
        prompt's hidden states, which a backend must either host in process or return."""
        return Requirements(
            generate=any_of(
                needs(Capability.IN_PROCESS_TORCH),
                needs(
                    Capability.HIDDEN_CAPTURE,
                    kinds=CaptureKinds(
                        kinds=frozenset({"residual"}),
                        locations=frozenset({"layer_input"}),
                        modes=frozenset({"all_tokens"}),
                    ),
                    hint=(
                        "the probe pass needs hidden-state capture, which this backend does not "
                        "return; run on huggingface or the offline vLLM engine"
                    ),
                ),
            ),
        )

    def steer_access(self) -> ModelAccess:
        """`ModelAccess.CAPTURE` when `probes` is a `ProbeSetFit` (the fit extracts hidden
        states), `ModelAccess.FACTS` for a fitted `ProbeSet` (identity checks read the session
        layout)."""
        if isinstance(self.probes, ProbeSetFit):
            return ModelAccess.CAPTURE
        return ModelAccess.FACTS

    def steer_fits(self) -> tuple[tuple[str, str], ...]:
        """One `("ProbeSetFit", "calibrated")` entry when the probes arrive as a fit recipe."""
        if isinstance(self.probes, ProbeSetFit):
            return (("ProbeSetFit", "calibrated"),)
        return ()

    def export_state(self) -> dict:
        """The fitted probe set under the `"probes"` key (after `steer()`)."""
        if self.probes is not None and not isinstance(self.probes, ProbeSetFit):
            return {"probes": self.probes}
        return {}

    def frozen_form(self, state: dict) -> tuple[str, dict]:
        """A same-class frozen form: the fitted `ProbeSet` plus the recipe's routing rules."""
        return "output_control/routed_decoding", {
            "probes": state["probes"],
            "rules": self.args.rules,
            "allow_model_mismatch": self.allow_model_mismatch,
        }

    def fit_identity(self):
        """The `ProbeSetFit` recipe, or None when the probes arrived fitted."""
        if isinstance(self.args.probes, ProbeSetFit):
            return self.args.probes
        return None

    def steer(
        self,
        model: PreTrainedModel | None = None,
        tokenizer: PreTrainedTokenizerBase | None = None,
        session: SteeringSession | None = None,
        **__,
    ) -> PreTrainedModel | None:
        """Attach the tokenizer, resolve the probes, and validate their identity.

        A `ProbeSetFit` is fitted here, on the model or session the pipeline provides (its
        `StatsSpec`, when present, is estimated first). A fitted `ProbeSet` is checked
        instead, venue-matched: with a live model, any probe whose recorded
        `model_fingerprint` differs raises unless `allow_model_mismatch=True` (probes with no
        recorded fingerprint are exempt); without one, the recorded `model_ref` and
        `model_type` are compared against the session layout's, with unrecorded values
        exempt.

        Args:
            model: The pipeline's model, or None on backends without a live model.
            tokenizer: Tokenizer used for splicing and padding. If None, attempts to retrieve
                from model attributes.
            session: `SteeringSession` scoped to this control, provided by the pipeline.

        Returns:
            The input model, unchanged.

        Raises:
            ValueError: If a fitted set's recorded identity differs from the venue's, or a
                route references a probe name the set does not define.
        """
        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)

        layout = None
        if model is None and session is not None:
            layout = session.layout

        if isinstance(self.probes, ProbeSetFit):
            self.probes = self.probes.fit(model, self.tokenizer, session=session)
        elif not self.allow_model_mismatch:
            live_fingerprint = None
            if model is not None:
                live_fingerprint = model_fingerprint(model)
            elif layout is not None and getattr(session, "in_process", False):
                live_fingerprint = layout.model_fingerprint
            if live_fingerprint is not None:
                mismatched = [
                    name for name, probe in self.probes.probes.items()
                    if probe.meta.get("model_fingerprint") not in (None, live_fingerprint)
                ]
                if mismatched:
                    raise ValueError(
                        "ProbeSet was fitted on a different model than this pipeline produced. "
                        "Pass a ProbeSetFit for steer-time fitting on the pipeline's final model, "
                        "or set allow_model_mismatch=True."
                    )
            elif layout is not None and layout.model_ref is not None:
                mismatched = [
                    name for name, probe in self.probes.probes.items()
                    if probe.meta.get("model_ref") not in (None, layout.model_ref)
                ]
                if mismatched:
                    raise ValueError(
                        "ProbeSet records a model reference that differs from the one this "
                        f"backend serves ({layout.model_ref!r}). Pass a ProbeSetFit for "
                        "steer-time fitting on the pipeline's final model, or set "
                        "allow_model_mismatch=True."
                    )

        if model is not None or (layout is not None and getattr(session, "in_process", False)):
            live_model_type = (
                getattr(model.config, "model_type", "unknown") if model is not None
                else (layout.model_type or "unknown")
            )
            if self.probes.model_type != live_model_type:
                raise ValueError(
                    f"ProbeSet was fitted on model_type {self.probes.model_type!r} but the "
                    f"pipeline's model is {live_model_type!r}."
                )
        elif (
            layout is not None
            and layout.model_type is not None
            and self.probes.model_type not in ("unknown", layout.model_type)
        ):
            raise ValueError(
                f"ProbeSet was fitted on model_type {self.probes.model_type!r} but this "
                f"backend serves {layout.model_type!r}."
            )

        self.rules.validate_names(set(self.probes.names))
        return model

    def decode(self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None,
               model: PreTrainedModel | None, logits_processors: LogitsProcessorList,
               stopping_criteria: StoppingCriteriaList, runtime_kwargs: dict | None,
               session: SteeringSession | None = None, **gen_kwargs) -> torch.Tensor:
        """Read the probes on the prompt, route each row, and execute the routed phase plans.

        Args:
            input_ids: Prompt token ids `[B, T]` (a 1-D tensor is treated as one row).
            attention_mask: Prompt attention mask matching `input_ids`, or None.
            model: The pipeline's model.
            logits_processors: Composed logits-processor stack, applied in every generated
                phase.
            stopping_criteria: Composed stopping-criteria stack, applied in every generated
                phase.
            runtime_kwargs: Per-call parameters (see the class docstring).

        Returns:
            Full sequence ids `[B, L]` (prompt + continuation), padded per row.

        Raises:
            RuntimeError: If no tokenizer is attached; `steer()` must run first.
            TypeError: If a matched action cannot be lowered to a phase plan.
            ValueError: If a matched plan contains a prompt-replacing `Fixed` phase
                (`replace=True`).

        Warns:
            UserWarning: If `canned_responses` carries keys that do not name a
                `Respond`/`Prefix` route.
        """
        if self.tokenizer is None:
            raise RuntimeError("RoutedDecoding requires a tokenizer; steer() must run first.")

        runtime_kwargs = runtime_kwargs or {}
        base_generate = resolve_generate_callable(model, runtime_kwargs, session)
        overrides = runtime_kwargs.get("canned_responses") or {}

        if input_ids.dim() == 1:
            input_ids = input_ids.unsqueeze(0)
        if attention_mask is not None and attention_mask.dim() == 1:
            attention_mask = attention_mask.unsqueeze(0)
        batch_size = input_ids.size(0)

        readings = self.probes.read(model, input_ids, attention_mask, session=session)
        matched = self.rules.route(readings.decisions)
        self.latest_routes = [route.name if route is not None else "default" for route in matched]

        if overrides:
            routes_by_name = {route.name: route for route in self.rules.routes}
            unusable = [
                key for key in overrides
                if key not in routes_by_name
                or not isinstance(routes_by_name[key].action, (Respond, Prefix))
            ]
            if unusable:
                warnings.warn(
                    f"canned_responses keys {unusable} do not name a Respond/Prefix route and "
                    "are ignored.",
                    UserWarning,
                )

        prompts = self.tokenizer.decode(input_ids, skip_special_tokens=True)

        final_sequences: list[torch.Tensor] = []
        for i in range(batch_size):
            route = matched[i]
            if route is not None:
                action = route.action
                if route.name in overrides and isinstance(action, (Respond, Prefix)):
                    action = replace(action, text=overrides[route.name])
            else:
                action = self.rules.default_action if self.rules.default_action is not None else Generate()
            plan = self._lower(action)

            row_full = input_ids[i:i + 1]
            if attention_mask is not None:
                row_ids = row_full[:, attention_mask[i].bool()]  # strip pad positions
            else:
                row_ids = row_full
            full = self._run_plan(
                plan, row_ids, prompts[i], {}, base_generate, False,
                logits_processors, stopping_criteria, gen_kwargs,
            )
            continuation = full[0][row_ids.size(1):]
            final_sequences.append(torch.cat([row_full[0], continuation]))

        # right padding keeps each prompt at offsets [0, T), which the pipeline's
        # prompt-length slicing requires regardless of the tokenizer's configured side
        padded = self.tokenizer.pad(
            {"input_ids": [seq.tolist() for seq in final_sequences]},
            padding=True,
            padding_side="right",
            return_tensors="pt",
        ).to(input_ids.device)
        return padded["input_ids"]

    @staticmethod
    def _lower(action) -> list:
        """Lower an action to a phase plan.

        Args:
            action: A `Respond`/`Prefix`/`Generate` (or any object with a `plan()` method), or
                a raw list/tuple of `Fixed`/`Generated` phases.

        Returns:
            The phase plan list.

        Raises:
            TypeError: If the action is neither plan-bearing nor a phase sequence.
            ValueError: If the plan contains a prompt-replacing `Fixed` phase (`replace=True`);
                routed actions append to the row's prompt, which the returned-continuation
                accounting relies on.
        """
        if hasattr(action, "plan"):
            plan = action.plan()
        elif isinstance(action, (list, tuple)):
            plan = list(action)
        else:
            raise TypeError(
                f"Cannot lower action of type {type(action).__name__} to a phase plan; use "
                "respond()/prefix()/generate() or a list of Fixed/Generated phases."
            )
        if any(isinstance(phase, Fixed) and phase.replace for phase in plan):
            raise ValueError(
                "RoutedDecoding does not support prompt-replacing Fixed phases (replace=True); "
                "routed actions append to the row's prompt."
            )
        return plan

Args = RoutedDecodingArgs class-attribute instance-attribute

RUNTIME_KWARGS_SCHEMA = [{'name': 'canned_responses', 'type': 'dict[str, str]', 'scope': 'call', 'description': 'Per-call override of Respond/Prefix text, keyed by route name.'}] class-attribute instance-attribute

args = self.Args.validate(*args, **kwargs) instance-attribute

enabled = True class-attribute instance-attribute

extract_after = extract_after instance-attribute

include_in_scoring = True class-attribute instance-attribute

same_model_forwards = True class-attribute instance-attribute

supports_batching = True class-attribute instance-attribute

tokenizer = None class-attribute instance-attribute

cleanup()

Release resources allocated during steer().

Override in subclasses that allocate GPU memory or other resources during steering. Default no-op.

Source code in steerability/algorithms/core/base_control.py

def cleanup(self) -> None:
    """Release resources allocated during `steer()`.

    Override in subclasses that allocate GPU memory or other resources during steering. Default
    no-op.
    """
    pass

clone_for_call(seed=None)

A configuration-preserving shallow clone for one generation call.

The clone shares steer-time artifacts (memories, steering vectors, attached tokenizers) with the original but has its own attribute namespace, so per-call attribute mutation on the clone never races another call using the original. When seed is given and the control defines reseed(seed), the clone's client-side RNG is re-seeded.

Parameters:

Name	Type	Description	Default
seed	int | None	Optional seed forwarded to the clone's reseed().	None

Returns:

Type	Description
BaseControl	The clone.

Source code in steerability/algorithms/core/base_control.py

def clone_for_call(self, seed: int | None = None) -> "BaseControl":
    """A configuration-preserving shallow clone for one generation call.

    The clone shares steer-time artifacts (memories, steering vectors, attached tokenizers)
    with the original but has its own attribute namespace, so per-call attribute mutation on
    the clone never races another call using the original. When `seed` is given and the
    control defines `reseed(seed)`, the clone's client-side RNG is re-seeded.

    Args:
        seed: Optional seed forwarded to the clone's `reseed()`.

    Returns:
        The clone.
    """
    clone = copy.copy(self)
    if seed is not None:
        reseed = getattr(clone, "reseed", None)
        if callable(reseed):
            reseed(seed)
    return clone

decode(input_ids, attention_mask, model, logits_processors, stopping_criteria, runtime_kwargs, session=None, **gen_kwargs)

Read the probes on the prompt, route each row, and execute the routed phase plans.

Parameters:

Name	Type	Description	Default
input_ids	Tensor	Prompt token ids [B, T] (a 1-D tensor is treated as one row).	required
attention_mask	Tensor | None	Prompt attention mask matching input_ids, or None.	required
model	PreTrainedModel | None	The pipeline's model.	required
logits_processors	LogitsProcessorList	Composed logits-processor stack, applied in every generated phase.	required
stopping_criteria	StoppingCriteriaList	Composed stopping-criteria stack, applied in every generated phase.	required
runtime_kwargs	dict | None	Per-call parameters (see the class docstring).	required

Returns:

Type	Description
Tensor	Full sequence ids [B, L] (prompt + continuation), padded per row.

Raises:

Type	Description
RuntimeError	If no tokenizer is attached; steer() must run first.
TypeError	If a matched action cannot be lowered to a phase plan.
ValueError	If a matched plan contains a prompt-replacing Fixed phase (replace=True).

Warns:

Type	Description
UserWarning	If canned_responses carries keys that do not name a Respond/Prefix route.

Source code in steerability/algorithms/output_control/routed_decoding/control.py

def decode(self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None,
           model: PreTrainedModel | None, logits_processors: LogitsProcessorList,
           stopping_criteria: StoppingCriteriaList, runtime_kwargs: dict | None,
           session: SteeringSession | None = None, **gen_kwargs) -> torch.Tensor:
    """Read the probes on the prompt, route each row, and execute the routed phase plans.

    Args:
        input_ids: Prompt token ids `[B, T]` (a 1-D tensor is treated as one row).
        attention_mask: Prompt attention mask matching `input_ids`, or None.
        model: The pipeline's model.
        logits_processors: Composed logits-processor stack, applied in every generated
            phase.
        stopping_criteria: Composed stopping-criteria stack, applied in every generated
            phase.
        runtime_kwargs: Per-call parameters (see the class docstring).

    Returns:
        Full sequence ids `[B, L]` (prompt + continuation), padded per row.

    Raises:
        RuntimeError: If no tokenizer is attached; `steer()` must run first.
        TypeError: If a matched action cannot be lowered to a phase plan.
        ValueError: If a matched plan contains a prompt-replacing `Fixed` phase
            (`replace=True`).

    Warns:
        UserWarning: If `canned_responses` carries keys that do not name a
            `Respond`/`Prefix` route.
    """
    if self.tokenizer is None:
        raise RuntimeError("RoutedDecoding requires a tokenizer; steer() must run first.")

    runtime_kwargs = runtime_kwargs or {}
    base_generate = resolve_generate_callable(model, runtime_kwargs, session)
    overrides = runtime_kwargs.get("canned_responses") or {}

    if input_ids.dim() == 1:
        input_ids = input_ids.unsqueeze(0)
    if attention_mask is not None and attention_mask.dim() == 1:
        attention_mask = attention_mask.unsqueeze(0)
    batch_size = input_ids.size(0)

    readings = self.probes.read(model, input_ids, attention_mask, session=session)
    matched = self.rules.route(readings.decisions)
    self.latest_routes = [route.name if route is not None else "default" for route in matched]

    if overrides:
        routes_by_name = {route.name: route for route in self.rules.routes}
        unusable = [
            key for key in overrides
            if key not in routes_by_name
            or not isinstance(routes_by_name[key].action, (Respond, Prefix))
        ]
        if unusable:
            warnings.warn(
                f"canned_responses keys {unusable} do not name a Respond/Prefix route and "
                "are ignored.",
                UserWarning,
            )

    prompts = self.tokenizer.decode(input_ids, skip_special_tokens=True)

    final_sequences: list[torch.Tensor] = []
    for i in range(batch_size):
        route = matched[i]
        if route is not None:
            action = route.action
            if route.name in overrides and isinstance(action, (Respond, Prefix)):
                action = replace(action, text=overrides[route.name])
        else:
            action = self.rules.default_action if self.rules.default_action is not None else Generate()
        plan = self._lower(action)

        row_full = input_ids[i:i + 1]
        if attention_mask is not None:
            row_ids = row_full[:, attention_mask[i].bool()]  # strip pad positions
        else:
            row_ids = row_full
        full = self._run_plan(
            plan, row_ids, prompts[i], {}, base_generate, False,
            logits_processors, stopping_criteria, gen_kwargs,
        )
        continuation = full[0][row_ids.size(1):]
        final_sequences.append(torch.cat([row_full[0], continuation]))

    # right padding keeps each prompt at offsets [0, T), which the pipeline's
    # prompt-length slicing requires regardless of the tokenizer's configured side
    padded = self.tokenizer.pad(
        {"input_ids": [seq.tolist() for seq in final_sequences]},
        padding=True,
        padding_side="right",
        return_tensors="pt",
    ).to(input_ids.device)
    return padded["input_ids"]

export_constraint(runtime_kwargs=None)

The control's declarative constrained-decoding source, or None.

A control whose per-step masking compiles from a declarative source returns a ConstraintSource; on a backend advertising Capability.GUIDED_DECODING the pipeline renders it onto the engine's native structured-output parameters in place of the control's live processor. The default returns None, which keeps the control on the live processor mechanism.

Parameters:

Name	Type	Description	Default
runtime_kwargs	dict | None	Per-call parameters supplied to generate().	None

Returns:

Type	Description
ConstraintSource | None	The constraint source, or None.

Source code in steerability/algorithms/output_control/base.py

def export_constraint(self, runtime_kwargs: dict | None = None) -> ConstraintSource | None:
    """The control's declarative constrained-decoding source, or None.

    A control whose per-step masking compiles from a declarative source returns a
    `ConstraintSource`; on a backend advertising `Capability.GUIDED_DECODING` the pipeline
    renders it onto the engine's native structured-output parameters in place of the
    control's live processor. The default returns None, which keeps the control on the live
    processor mechanism.

    Args:
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        The constraint source, or None.
    """
    return None

export_generation_params(runtime_kwargs=None)

The control's sampling-expressible contribution, or None.

A control whose behavior is expressible as normalized generation parameters returns a mapping over a subset of stop_strings, stop_token_ids, max_new_tokens, and min_new_tokens; the pipeline merges it into the call's GenerationParams (stop rules union with the caller's; token bounds only tighten) and does not additionally collect the control's live processors and criteria for that call, so the control executes on every backend through the session's composed stop rules. The default returns None, which keeps the control on the live processor/criteria mechanism.

Parameters:

Name	Type	Description	Default
runtime_kwargs	dict | None	Per-call parameters supplied to generate().	None

Returns:

Type	Description
Mapping[str, Any] | None	The parameter contribution, or None.

Source code in steerability/algorithms/output_control/base.py

def export_generation_params(self, runtime_kwargs: dict | None = None) -> Mapping[str, Any] | None:
    """The control's sampling-expressible contribution, or None.

    A control whose behavior is expressible as normalized generation parameters returns a
    mapping over a subset of `stop_strings`, `stop_token_ids`, `max_new_tokens`, and
    `min_new_tokens`; the pipeline merges it into the call's `GenerationParams` (stop rules
    union with the caller's; token bounds only tighten) and does not additionally collect
    the control's live processors and criteria for that call, so the control executes on
    every backend through the session's composed stop rules. The default returns None, which
    keeps the control on the live processor/criteria mechanism.

    Args:
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        The parameter contribution, or None.
    """
    return None

export_processor_spec(runtime_kwargs=None)

The control's engine-hosted processor form, or None.

A control whose per-step logit math is expressible in an engine's served processor vocabulary returns a ProcessorSpec; on a backend advertising Capability.PER_STEP_LOGIT_SPECS with the spec's kind, the pipeline submits it as a ProcessorSpecEntry in place of the control's live processor. The default returns None, which keeps the control on the live processor mechanism.

Parameters:

Name	Type	Description	Default
runtime_kwargs	dict | None	Per-call parameters supplied to generate().	None

Returns:

Type	Description
ProcessorSpec | None	The processor spec, or None.

Source code in steerability/algorithms/output_control/base.py

def export_processor_spec(self, runtime_kwargs: dict | None = None) -> ProcessorSpec | None:
    """The control's engine-hosted processor form, or None.

    A control whose per-step logit math is expressible in an engine's served processor
    vocabulary returns a `ProcessorSpec`; on a backend advertising
    `Capability.PER_STEP_LOGIT_SPECS` with the spec's kind, the pipeline submits it as a
    `ProcessorSpecEntry` in place of the control's live processor. The default returns
    None, which keeps the control on the live processor mechanism.

    Args:
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        The processor spec, or None.
    """
    return None

export_state()

The fitted probe set under the "probes" key (after steer()).

Source code in steerability/algorithms/output_control/routed_decoding/control.py

def export_state(self) -> dict:
    """The fitted probe set under the `"probes"` key (after `steer()`)."""
    if self.probes is not None and not isinstance(self.probes, ProbeSetFit):
        return {"probes": self.probes}
    return {}

fit_identity()

The ProbeSetFit recipe, or None when the probes arrived fitted.

Source code in steerability/algorithms/output_control/routed_decoding/control.py

def fit_identity(self):
    """The `ProbeSetFit` recipe, or None when the probes arrived fitted."""
    if isinstance(self.args.probes, ProbeSetFit):
        return self.args.probes
    return None

frozen_form(state)

A same-class frozen form: the fitted ProbeSet plus the recipe's routing rules.

Source code in steerability/algorithms/output_control/routed_decoding/control.py

def frozen_form(self, state: dict) -> tuple[str, dict]:
    """A same-class frozen form: the fitted `ProbeSet` plus the recipe's routing rules."""
    return "output_control/routed_decoding", {
        "probes": state["probes"],
        "rules": self.args.rules,
        "allow_model_mismatch": self.allow_model_mismatch,
    }

get_logits_processors(input_ids, runtime_kwargs, **kwargs)

The control's logits processors for the current generation.

Called once per generate() / compute_logprobs() call, after input and state controls have prepared the prompt (mirrors StateControl.get_hooks). **kwargs carries attention_mask and the caller's generation kwargs. Returned objects follow the HF LogitsProcessor convention; in-list order is preserved by the composition.

A processor must behave as a function of (prefix_ids, scores). Internal state is permitted only as memoization keyed on the prefix and must re-derive on a prefix mismatch, since drivers may restart, rewind, or reorder sequences, and scoring replays prefixes teacher-forced (subclass common.processors.base.PrefixKeyedProcessor to satisfy this mechanically). Return fresh processor instances from this hook; it is invoked once per call precisely so that per-generation state is isolated.

Parameters:

Name	Type	Description	Default
input_ids	Tensor	The steered prompt token ids [batch, seq_len].	required
runtime_kwargs	dict | None	Per-call parameters supplied to generate().	required

Returns:

Type	Description
list	A list of HF LogitsProcessor-style objects.

Source code in steerability/algorithms/output_control/base.py

def get_logits_processors(self, input_ids: torch.Tensor, runtime_kwargs: dict | None, **kwargs) -> list:
    """The control's logits processors for the current generation.

    Called once per `generate()` / `compute_logprobs()` call, after input and state
    controls have prepared the prompt (mirrors `StateControl.get_hooks`). `**kwargs`
    carries `attention_mask` and the caller's generation kwargs. Returned objects
    follow the HF `LogitsProcessor` convention; in-list order is preserved by the composition.

    A processor must behave as a function of `(prefix_ids, scores)`. Internal state is
    permitted only as memoization keyed on the prefix and must re-derive on a prefix mismatch,
    since drivers may restart, rewind, or reorder sequences, and scoring replays prefixes
    teacher-forced (subclass `common.processors.base.PrefixKeyedProcessor` to satisfy this
    mechanically). Return fresh processor instances from this hook; it is invoked once per call
    precisely so that per-generation state is isolated.

    Args:
        input_ids: The steered prompt token ids `[batch, seq_len]`.
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        A list of HF `LogitsProcessor`-style objects.
    """
    return []

get_stopping_criteria(input_ids, runtime_kwargs, **kwargs)

The control's stopping criteria.

Not applied during scoring (there is no loop to stop). Same call convention as get_logits_processors.

Parameters:

Name	Type	Description	Default
input_ids	Tensor	The steered prompt token ids [batch, seq_len].	required
runtime_kwargs	dict | None	Per-call parameters supplied to generate().	required

Returns:

Type	Description
list	A list of HF StoppingCriteria-style objects.

Source code in steerability/algorithms/output_control/base.py

def get_stopping_criteria(self, input_ids: torch.Tensor, runtime_kwargs: dict | None, **kwargs) -> list:
    """The control's stopping criteria.

    Not applied during scoring (there is no loop to stop). Same call convention as
    `get_logits_processors`.

    Args:
        input_ids: The steered prompt token ids `[batch, seq_len]`.
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        A list of HF `StoppingCriteria`-style objects.
    """
    return []

max_rollouts_per_query()

1: every route lowers to at most one Generated phase, and the probe pass is a read-only capture rather than a generation.

Source code in steerability/algorithms/output_control/routed_decoding/control.py

def max_rollouts_per_query(self) -> int:
    """1: every route lowers to at most one `Generated` phase, and the probe pass is a
    read-only capture rather than a generation."""
    return 1

plan(prompt_text, params)

Return the phase plan for one example. Subclasses override.

Source code in steerability/algorithms/output_control/common/drivers/phased.py

def plan(self, prompt_text: str, params: dict) -> list:
    """Return the phase plan for one example. Subclasses override."""
    raise NotImplementedError

requirements()

In-process torch or hidden-state capture at generate; the probe pass reads the prompt's hidden states, which a backend must either host in process or return.

Source code in steerability/algorithms/output_control/routed_decoding/control.py

def requirements(self) -> Requirements:
    """In-process torch or hidden-state capture at generate; the probe pass reads the
    prompt's hidden states, which a backend must either host in process or return."""
    return Requirements(
        generate=any_of(
            needs(Capability.IN_PROCESS_TORCH),
            needs(
                Capability.HIDDEN_CAPTURE,
                kinds=CaptureKinds(
                    kinds=frozenset({"residual"}),
                    locations=frozenset({"layer_input"}),
                    modes=frozenset({"all_tokens"}),
                ),
                hint=(
                    "the probe pass needs hidden-state capture, which this backend does not "
                    "return; run on huggingface or the offline vLLM engine"
                ),
            ),
        ),
    )

steer(model=None, tokenizer=None, session=None, **__)

Attach the tokenizer, resolve the probes, and validate their identity.

A ProbeSetFit is fitted here, on the model or session the pipeline provides (its StatsSpec, when present, is estimated first). A fitted ProbeSet is checked instead, venue-matched: with a live model, any probe whose recorded model_fingerprint differs raises unless allow_model_mismatch=True (probes with no recorded fingerprint are exempt); without one, the recorded model_ref and model_type are compared against the session layout's, with unrecorded values exempt.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	The pipeline's model, or None on backends without a live model.	None
tokenizer	PreTrainedTokenizerBase | None	Tokenizer used for splicing and padding. If None, attempts to retrieve from model attributes.	None
session	SteeringSession | None	SteeringSession scoped to this control, provided by the pipeline.	None

Returns:

Type	Description
PreTrainedModel | None	The input model, unchanged.

Raises:

Type	Description
ValueError	If a fitted set's recorded identity differs from the venue's, or a route references a probe name the set does not define.

Source code in steerability/algorithms/output_control/routed_decoding/control.py

def steer(
    self,
    model: PreTrainedModel | None = None,
    tokenizer: PreTrainedTokenizerBase | None = None,
    session: SteeringSession | None = None,
    **__,
) -> PreTrainedModel | None:
    """Attach the tokenizer, resolve the probes, and validate their identity.

    A `ProbeSetFit` is fitted here, on the model or session the pipeline provides (its
    `StatsSpec`, when present, is estimated first). A fitted `ProbeSet` is checked
    instead, venue-matched: with a live model, any probe whose recorded
    `model_fingerprint` differs raises unless `allow_model_mismatch=True` (probes with no
    recorded fingerprint are exempt); without one, the recorded `model_ref` and
    `model_type` are compared against the session layout's, with unrecorded values
    exempt.

    Args:
        model: The pipeline's model, or None on backends without a live model.
        tokenizer: Tokenizer used for splicing and padding. If None, attempts to retrieve
            from model attributes.
        session: `SteeringSession` scoped to this control, provided by the pipeline.

    Returns:
        The input model, unchanged.

    Raises:
        ValueError: If a fitted set's recorded identity differs from the venue's, or a
            route references a probe name the set does not define.
    """
    self.tokenizer = tokenizer or getattr(model, "tokenizer", None)

    layout = None
    if model is None and session is not None:
        layout = session.layout

    if isinstance(self.probes, ProbeSetFit):
        self.probes = self.probes.fit(model, self.tokenizer, session=session)
    elif not self.allow_model_mismatch:
        live_fingerprint = None
        if model is not None:
            live_fingerprint = model_fingerprint(model)
        elif layout is not None and getattr(session, "in_process", False):
            live_fingerprint = layout.model_fingerprint
        if live_fingerprint is not None:
            mismatched = [
                name for name, probe in self.probes.probes.items()
                if probe.meta.get("model_fingerprint") not in (None, live_fingerprint)
            ]
            if mismatched:
                raise ValueError(
                    "ProbeSet was fitted on a different model than this pipeline produced. "
                    "Pass a ProbeSetFit for steer-time fitting on the pipeline's final model, "
                    "or set allow_model_mismatch=True."
                )
        elif layout is not None and layout.model_ref is not None:
            mismatched = [
                name for name, probe in self.probes.probes.items()
                if probe.meta.get("model_ref") not in (None, layout.model_ref)
            ]
            if mismatched:
                raise ValueError(
                    "ProbeSet records a model reference that differs from the one this "
                    f"backend serves ({layout.model_ref!r}). Pass a ProbeSetFit for "
                    "steer-time fitting on the pipeline's final model, or set "
                    "allow_model_mismatch=True."
                )

    if model is not None or (layout is not None and getattr(session, "in_process", False)):
        live_model_type = (
            getattr(model.config, "model_type", "unknown") if model is not None
            else (layout.model_type or "unknown")
        )
        if self.probes.model_type != live_model_type:
            raise ValueError(
                f"ProbeSet was fitted on model_type {self.probes.model_type!r} but the "
                f"pipeline's model is {live_model_type!r}."
            )
    elif (
        layout is not None
        and layout.model_type is not None
        and self.probes.model_type not in ("unknown", layout.model_type)
    ):
        raise ValueError(
            f"ProbeSet was fitted on model_type {self.probes.model_type!r} but this "
            f"backend serves {layout.model_type!r}."
        )

    self.rules.validate_names(set(self.probes.names))
    return model

steer_access()

ModelAccess.CAPTURE when probes is a ProbeSetFit (the fit extracts hidden states), ModelAccess.FACTS for a fitted ProbeSet (identity checks read the session layout).

Source code in steerability/algorithms/output_control/routed_decoding/control.py

def steer_access(self) -> ModelAccess:
    """`ModelAccess.CAPTURE` when `probes` is a `ProbeSetFit` (the fit extracts hidden
    states), `ModelAccess.FACTS` for a fitted `ProbeSet` (identity checks read the session
    layout)."""
    if isinstance(self.probes, ProbeSetFit):
        return ModelAccess.CAPTURE
    return ModelAccess.FACTS

steer_fits()

One ("ProbeSetFit", "calibrated") entry when the probes arrive as a fit recipe.

Source code in steerability/algorithms/output_control/routed_decoding/control.py

def steer_fits(self) -> tuple[tuple[str, str], ...]:
    """One `("ProbeSetFit", "calibrated")` entry when the probes arrive as a fit recipe."""
    if isinstance(self.probes, ProbeSetFit):
        return (("ProbeSetFit", "calibrated"),)
    return ()

routing

Predicates and ordered routes over named per-row boolean decisions.

A Predicate is a boolean expression over decision names, built from P(name) leaves and the operators & (and), | (or), and ~ (not), and evaluated per row against a mapping from decision name to per-row values. Probes are the canonical producer of such decisions, but any named boolean source works. Router holds an ordered list of Routes with first-match-wins semantics, evaluated independently per row, so one batched call can route each prompt to a different route.

Actions are opaque payloads: a Route carries whatever the consumer interprets (e.g. a decoding driver lowers actions to phase plans). This module depends only on torch and the standard library.

Predicate

Bases: ABC

Boolean expression over named per-row decisions, evaluated per row.

Leaves are created with P(name); composites are built with the operators & (and), | (or), and ~ (not). evaluate() takes a mapping from decision name to per-row values (1-D bool tensors of a common length, or a bare bool for single-row batches) and returns a bool tensor of shape [num_rows]. repr() renders the expression infix, e.g. (legal & ~advice).

Source code in steerability/algorithms/output_control/routed_decoding/routing.py

class Predicate(ABC):
    """Boolean expression over named per-row decisions, evaluated per row.

    Leaves are created with `P(name)`; composites are built with the operators `&` (and),
    `|` (or), and `~` (not). `evaluate()` takes a mapping from decision name to per-row values
    (1-D bool tensors of a common length, or a bare bool for single-row batches) and returns a
    bool tensor of shape `[num_rows]`. `repr()` renders the expression infix, e.g.
    ``(legal & ~advice)``.
    """

    def evaluate(self, decisions: Mapping[str, torch.Tensor | bool]) -> torch.Tensor:
        """Evaluate the predicate against per-row decisions.

        Args:
            decisions: Mapping from decision name to a 1-D bool tensor of shape `[num_rows]`
                (a bare bool is accepted only when `num_rows == 1`).

        Returns:
            Bool tensor of shape `[num_rows]`; True where the predicate holds for that row.

        Raises:
            KeyError: If the predicate references a decision name absent from `decisions`; the
                message lists the available names.
            ValueError: If the decision tensors are malformed (see `evaluate` requirements).
        """
        normalized, _ = _normalize_decisions(decisions)
        return self._eval(normalized)

    @abstractmethod
    def _eval(self, decisions: dict[str, torch.Tensor]) -> torch.Tensor:
        """Evaluate against an already-normalized decisions mapping."""

    @abstractmethod
    def decision_names(self) -> set[str]:
        """The set of decision names this predicate references."""

    def __and__(self, other: "Predicate") -> "Predicate":
        if not isinstance(other, Predicate):
            return NotImplemented
        return _And(self, other)

    def __or__(self, other: "Predicate") -> "Predicate":
        if not isinstance(other, Predicate):
            return NotImplemented
        return _Or(self, other)

    def __invert__(self) -> "Predicate":
        return _Not(self)

decision_names() abstractmethod

The set of decision names this predicate references.

Source code in steerability/algorithms/output_control/routed_decoding/routing.py

@abstractmethod
def decision_names(self) -> set[str]:
    """The set of decision names this predicate references."""

evaluate(decisions)

Evaluate the predicate against per-row decisions.

Parameters:

Name	Type	Description	Default
decisions	Mapping[str, Tensor | bool]	Mapping from decision name to a 1-D bool tensor of shape [num_rows] (a bare bool is accepted only when num_rows == 1).	required

Returns:

Type	Description
Tensor	Bool tensor of shape [num_rows]; True where the predicate holds for that row.

Raises:

Type	Description
KeyError	If the predicate references a decision name absent from decisions; the message lists the available names.
ValueError	If the decision tensors are malformed (see evaluate requirements).

Source code in steerability/algorithms/output_control/routed_decoding/routing.py

def evaluate(self, decisions: Mapping[str, torch.Tensor | bool]) -> torch.Tensor:
    """Evaluate the predicate against per-row decisions.

    Args:
        decisions: Mapping from decision name to a 1-D bool tensor of shape `[num_rows]`
            (a bare bool is accepted only when `num_rows == 1`).

    Returns:
        Bool tensor of shape `[num_rows]`; True where the predicate holds for that row.

    Raises:
        KeyError: If the predicate references a decision name absent from `decisions`; the
            message lists the available names.
        ValueError: If the decision tensors are malformed (see `evaluate` requirements).
    """
    normalized, _ = _normalize_decisions(decisions)
    return self._eval(normalized)

Route dataclass

One named route: a predicate and the action to take when it matches.

Attributes:

Name	Type	Description
name	str	Route name, unique within a Router; keys diagnostics and per-call action overrides.
when	Predicate	The predicate that must hold for a row to match this route.
action	Any	Opaque payload interpreted by the consumer (e.g. lowered to a phase plan by a decoding driver).

Source code in steerability/algorithms/output_control/routed_decoding/routing.py

@dataclass(frozen=True)
class Route:
    """One named route: a predicate and the action to take when it matches.

    Attributes:
        name: Route name, unique within a `Router`; keys diagnostics and per-call action
            overrides.
        when: The predicate that must hold for a row to match this route.
        action: Opaque payload interpreted by the consumer (e.g. lowered to a phase plan by a
            decoding driver).
    """

    name: str
    when: Predicate
    action: Any

    def __post_init__(self):
        if not isinstance(self.name, str) or not self.name:
            raise ValueError(f"Route name must be a non-empty string; got {self.name!r}.")
        if not isinstance(self.when, Predicate):
            raise TypeError(
                f"Route '{self.name}': `when` must be a Predicate (build one with P(name) "
                f"and &, |, ~); got {type(self.when).__name__}."
            )

action instance-attribute

name instance-attribute

when instance-attribute

Router

Ordered routes with first-match-wins semantics, evaluated independently per row.

Each route pairs a Predicate with an action payload. route() evaluates every route's predicate over all rows, then assigns each row the first route whose predicate holds; rows matching no route fall to the default (returned as None, with default_action available to the consumer).

Parameters:

Name	Type	Description	Default
routes	Sequence[Route]	Ordered routes; earlier routes take precedence. Names must be unique.	required
default_action	Any	Action payload for rows matching no route. Opaque, like Route.action.	None

Raises:

Type	Description
ValueError	If two routes share a name.
TypeError	If an entry is not a Route.

Source code in steerability/algorithms/output_control/routed_decoding/routing.py

class Router:
    """Ordered routes with first-match-wins semantics, evaluated independently per row.

    Each route pairs a `Predicate` with an action payload. `route()` evaluates every route's
    predicate over all rows, then assigns each row the first route whose predicate holds; rows
    matching no route fall to the default (returned as None, with `default_action` available to
    the consumer).

    Args:
        routes: Ordered routes; earlier routes take precedence. Names must be unique.
        default_action: Action payload for rows matching no route. Opaque, like `Route.action`.

    Raises:
        ValueError: If two routes share a name.
        TypeError: If an entry is not a `Route`.
    """

    def __init__(self, routes: Sequence[Route], default_action: Any = None):
        routes = tuple(routes)
        for entry in routes:
            if not isinstance(entry, Route):
                raise TypeError(f"Router entries must be Route instances; got {type(entry).__name__}.")
        seen: set[str] = set()
        for entry in routes:
            if entry.name in seen:
                raise ValueError(f"Duplicate route name '{entry.name}'; route names must be unique.")
            seen.add(entry.name)
        self.routes = routes
        self.default_action = default_action

    def route(self, decisions: Mapping[str, torch.Tensor | bool]) -> list[Route | None]:
        """Match each row to its first satisfied route.

        Every route's predicate is evaluated once over all rows, then each row takes the first
        route whose predicate holds for it. There is no batch-wide short-circuit; rows are
        routed independently.

        Args:
            decisions: Mapping from decision name to per-row decisions (see
                `Predicate.evaluate`).

        Returns:
            One entry per row: the matched `Route`, or None for rows matching no route (the
            default route).

        Raises:
            KeyError: If a route references a decision name absent from `decisions`.
            ValueError: If the decision tensors are malformed.
        """
        normalized, num_rows = _normalize_decisions(decisions)
        masks = [entry.when._eval(normalized) for entry in self.routes]
        matched: list[Route | None] = []
        for row in range(num_rows):
            match: Route | None = None
            for entry, mask in zip(self.routes, masks):
                if bool(mask[row]):
                    match = entry
                    break
            matched.append(match)
        return matched

    def decision_names(self) -> set[str]:
        """The union of decision names referenced by all routes."""
        names: set[str] = set()
        for entry in self.routes:
            names |= entry.when.decision_names()
        return names

    def validate_names(self, available: set[str]) -> None:
        """Check that every referenced decision name exists among `available`.

        Args:
            available: The decision names the consumer provides at routing time.

        Raises:
            ValueError: If any route references a decision name absent from `available`; the
                message names the missing decisions.
        """
        missing = sorted(self.decision_names() - set(available))
        if missing:
            raise ValueError(
                f"Routes reference unknown decision name(s) {missing}; available decisions: "
                f"{sorted(available)}."
            )

    def describe(self) -> str:
        """Render the route set as a plain-text flowchart.

        One line per route, in precedence order, followed by the default line. The action
        column uses `str(action)` when the action type defines its own `__str__`, else the
        type name.

        Returns:
            The multi-line flowchart string.
        """
        name_width = max((len(entry.name) for entry in self.routes), default=0)
        name_width = max(name_width, len("default"))
        pred_reprs = [repr(entry.when) for entry in self.routes]
        pred_width = max((len(p) for p in pred_reprs), default=0)
        index_width = len(str(len(self.routes))) if self.routes else 1
        body_width = index_width + 2 + name_width + 3 + 3 + pred_width + 3

        lines = ["Router"]
        for i, (entry, pred) in enumerate(zip(self.routes, pred_reprs), start=1):
            body = f"{i:>{index_width}}. {entry.name:<{name_width}}   if {pred:<{pred_width}}   "
            lines.append(f"├─ {body}-> {_action_label(entry.action)}")
        lines.append(f"└─ {'default':<{body_width}}-> {_action_label(self.default_action)}")
        return "\n".join(lines)

    def __repr__(self) -> str:
        route_names = ", ".join(entry.name for entry in self.routes)
        return f"Router([{route_names}], default_action={_action_label(self.default_action)})"

default_action = default_action instance-attribute

routes = routes instance-attribute

decision_names()

The union of decision names referenced by all routes.

Source code in steerability/algorithms/output_control/routed_decoding/routing.py

def decision_names(self) -> set[str]:
    """The union of decision names referenced by all routes."""
    names: set[str] = set()
    for entry in self.routes:
        names |= entry.when.decision_names()
    return names

describe()

Render the route set as a plain-text flowchart.

One line per route, in precedence order, followed by the default line. The action column uses str(action) when the action type defines its own __str__, else the type name.

Returns:

Type	Description
str	The multi-line flowchart string.

Source code in steerability/algorithms/output_control/routed_decoding/routing.py

def describe(self) -> str:
    """Render the route set as a plain-text flowchart.

    One line per route, in precedence order, followed by the default line. The action
    column uses `str(action)` when the action type defines its own `__str__`, else the
    type name.

    Returns:
        The multi-line flowchart string.
    """
    name_width = max((len(entry.name) for entry in self.routes), default=0)
    name_width = max(name_width, len("default"))
    pred_reprs = [repr(entry.when) for entry in self.routes]
    pred_width = max((len(p) for p in pred_reprs), default=0)
    index_width = len(str(len(self.routes))) if self.routes else 1
    body_width = index_width + 2 + name_width + 3 + 3 + pred_width + 3

    lines = ["Router"]
    for i, (entry, pred) in enumerate(zip(self.routes, pred_reprs), start=1):
        body = f"{i:>{index_width}}. {entry.name:<{name_width}}   if {pred:<{pred_width}}   "
        lines.append(f"├─ {body}-> {_action_label(entry.action)}")
    lines.append(f"└─ {'default':<{body_width}}-> {_action_label(self.default_action)}")
    return "\n".join(lines)

route(decisions)

Match each row to its first satisfied route.

Every route's predicate is evaluated once over all rows, then each row takes the first route whose predicate holds for it. There is no batch-wide short-circuit; rows are routed independently.

Parameters:

Name	Type	Description	Default
decisions	Mapping[str, Tensor | bool]	Mapping from decision name to per-row decisions (see Predicate.evaluate).	required

Returns:

Type	Description
list[Route | None]	One entry per row: the matched Route, or None for rows matching no route (the
list[Route | None]	default route).

Raises:

Type	Description
KeyError	If a route references a decision name absent from decisions.
ValueError	If the decision tensors are malformed.

Source code in steerability/algorithms/output_control/routed_decoding/routing.py

def route(self, decisions: Mapping[str, torch.Tensor | bool]) -> list[Route | None]:
    """Match each row to its first satisfied route.

    Every route's predicate is evaluated once over all rows, then each row takes the first
    route whose predicate holds for it. There is no batch-wide short-circuit; rows are
    routed independently.

    Args:
        decisions: Mapping from decision name to per-row decisions (see
            `Predicate.evaluate`).

    Returns:
        One entry per row: the matched `Route`, or None for rows matching no route (the
        default route).

    Raises:
        KeyError: If a route references a decision name absent from `decisions`.
        ValueError: If the decision tensors are malformed.
    """
    normalized, num_rows = _normalize_decisions(decisions)
    masks = [entry.when._eval(normalized) for entry in self.routes]
    matched: list[Route | None] = []
    for row in range(num_rows):
        match: Route | None = None
        for entry, mask in zip(self.routes, masks):
            if bool(mask[row]):
                match = entry
                break
        matched.append(match)
    return matched

validate_names(available)

Check that every referenced decision name exists among available.

Parameters:

Name	Type	Description	Default
available	set[str]	The decision names the consumer provides at routing time.	required

Raises:

Type	Description
ValueError	If any route references a decision name absent from available; the message names the missing decisions.

Source code in steerability/algorithms/output_control/routed_decoding/routing.py

def validate_names(self, available: set[str]) -> None:
    """Check that every referenced decision name exists among `available`.

    Args:
        available: The decision names the consumer provides at routing time.

    Raises:
        ValueError: If any route references a decision name absent from `available`; the
            message names the missing decisions.
    """
    missing = sorted(self.decision_names() - set(available))
    if missing:
        raise ValueError(
            f"Routes reference unknown decision name(s) {missing}; available decisions: "
            f"{sorted(available)}."
        )

P(name)

Leaf predicate over the named decision.

Parameters:

Name	Type	Description	Default
name	str	The decision name, matching a key of the decisions mapping at evaluation time.	required

Returns:

Type	Description
Predicate	A Predicate that reads the named per-row decision.

Source code in steerability/algorithms/output_control/routed_decoding/routing.py

def P(name: str) -> Predicate:
    """Leaf predicate over the named decision.

    Args:
        name: The decision name, matching a key of the decisions mapping at evaluation time.

    Returns:
        A `Predicate` that reads the named per-row decision.
    """
    return _Decision(name)