Common library

steerability.algorithms.state_control.common

State control component library.

estimators

Estimator components for state control.

base

Base class for estimators that learn artifacts during steer().

T = TypeVar('T') module-attribute

BaseEstimator

Bases: ABC, Generic[T]

Learns an artifact from a model and data during the steer() phase.

The type parameter T is the artifact type (e.g., SteeringVector).

Source code in steerability/algorithms/state_control/common/estimators/base.py

class BaseEstimator(ABC, Generic[T]):
    """Learns an artifact from a model and data during the steer() phase.

    The type parameter T is the artifact type (e.g., SteeringVector).
    """

    @abstractmethod
    def fit(
        self,
        model: PreTrainedModel,
        tokenizer: PreTrainedTokenizerBase,
        **kwargs,
    ) -> T:
        """Fit the estimator and return the learned artifact.

        Args:
            model: The model being steered. Used to extract hidden states.
            tokenizer: Tokenizer for encoding training data.

        Returns:
            The learned artifact of type T.
        """
        ...

fit(model, tokenizer, **kwargs) abstractmethod

Fit the estimator and return the learned artifact.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	The model being steered. Used to extract hidden states.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for encoding training data.	required

Returns:

Type	Description
T	The learned artifact of type T.

Source code in steerability/algorithms/state_control/common/estimators/base.py

@abstractmethod
def fit(
    self,
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    **kwargs,
) -> T:
    """Fit the estimator and return the learned artifact.

    Args:
        model: The model being steered. Used to extract hidden states.
        tokenizer: Tokenizer for encoding training data.

    Returns:
        The learned artifact of type T.
    """
    ...

contrastive_direction

Contrastive direction estimator using paired PCA.

PcaMethod = Literal['pca_pairwise', 'pca_center'] module-attribute

logger = logging.getLogger(__name__) module-attribute

ContrastiveDirectionEstimator

Bases: BaseEstimator[SteeringVector]

Learns per-layer direction vectors from contrastive text pairs via PCA.

Two PCA variants are supported, selected by spec.method:

• pca_pairwise: centers each pair (H_l^+, H_l^-) at its midpoint, giving the samples ±(H_l^+ - H_l^-)/2, and takes the first principal component of that symmetric set.
• pca_center: fits PCA on the union of positive and negative pooled activations centered by their grand mean: vector_l = PCA(H_l^+ - mu_l, H_l^- - mu_l) with mu_l the mean over all examples of both classes.

For both methods the first principal component is oriented so positive examples project above negative examples (see _orient_direction).

Examples are rendered via render_for_model according to spec.prompt_format and tokenized with add_special_tokens=False for chat-templated text.

Source code in steerability/algorithms/state_control/common/estimators/contrastive_direction.py

class ContrastiveDirectionEstimator(BaseEstimator[SteeringVector]):
    """Learns per-layer direction vectors from contrastive text pairs via PCA.

    Two PCA variants are supported, selected by `spec.method`:

    - `pca_pairwise`: centers each pair `(H_l^+, H_l^-)` at its midpoint, giving the samples
        `±(H_l^+ - H_l^-)/2`, and takes the first principal component of that symmetric set.
    - `pca_center`: fits PCA on the union of positive and negative pooled activations centered by
        their grand mean: `vector_l = PCA(H_l^+ - mu_l, H_l^- - mu_l)` with `mu_l` the mean over all
        examples of both classes.

    For both methods the first principal component is oriented so positive examples project above
    negative examples (see `_orient_direction`).

    Examples are rendered via `render_for_model` according to `spec.prompt_format` and tokenized with
    `add_special_tokens=False` for chat-templated text.
    """

    def fit(
        self,
        model: PreTrainedModel | None,
        tokenizer: PreTrainedTokenizerBase,
        *,
        data: ContrastivePairs,
        spec: VectorTrainSpec,
        on_progress: Callable[[int, int], None] | None = None,
        session: "SteeringSession | None" = None,
    ) -> SteeringVector:
        """Extract contrastive direction vectors.

        Args:
            model: Model to extract hidden states from, or None to extract through `session`.
            tokenizer: Tokenizer for encoding the contrastive pairs.
            data: The positive/negative text pairs.
            spec: Training configuration (method, accumulate, batch_size).
            on_progress: Optional `(completed, total)` callback fired as each forward-pass batch
                finishes. `total` covers both positive and negative passes.
            session: A `SteeringSession` serving hidden-state capture when no live model is
                available.

        Returns:
            SteeringVector with one direction per layer.
        """
        device = next(model.parameters()).device if model is not None else torch.device("cpu")
        if model is not None:
            model_type = getattr(model.config, "model_type", "unknown")
            session_meta: dict = {}
        else:
            model_type, session_meta = session_artifact_identity(session)

        # render full texts according to prompt_format (shared with inference)
        rendered = render_contrastive(tokenizer, data, spec.prompt_format)

        logger.debug(
            "Tokenizing %d positive and %d negative examples", len(rendered.pos_texts), len(rendered.neg_texts)
        )

        # tokenize
        enc_pos = tokenize_texts(tokenizer, rendered.pos_texts, device, add_special_tokens=rendered.add_special_tokens)
        enc_neg = tokenize_texts(tokenizer, rendered.neg_texts, device, add_special_tokens=rendered.add_special_tokens)

        # tokenize prompts separately if needed for suffix-only
        prompt_enc = None
        if spec.accumulate == "suffix-only" and rendered.prompt_texts is not None:
            prompt_enc = tokenize_texts(
                tokenizer, rendered.prompt_texts, device, add_special_tokens=rendered.add_special_tokens
            )
            prompt_enc = {k: v.cpu() for k, v in prompt_enc.items()}

        # extract hidden states
        logger.debug("Extracting hidden states with batch_size=%d", spec.batch_size)
        n_pos = enc_pos["input_ids"].size(0)
        n_neg = enc_neg["input_ids"].size(0)
        total_batches = math.ceil(n_pos / spec.batch_size) + math.ceil(n_neg / spec.batch_size)
        completed = {"n": 0}

        def _tick() -> None:
            completed["n"] += 1
            if on_progress is not None:
                on_progress(completed["n"], total_batches)

        if on_progress is not None:
            on_progress(0, total_batches)
        hs_pos, mask_pos = capture_hidden(
            enc_pos, model=model, session=session,
            batch_size=spec.batch_size, on_batch=_tick, location=spec.location,
        )
        hs_neg, mask_neg = capture_hidden(
            enc_neg, model=model, session=session,
            batch_size=spec.batch_size, on_batch=_tick, location=spec.location,
        )

        # select spans against the returned layout (spans are mask-derived, so the returned
        # mask keeps them aligned when a remote capture re-packs rows)
        def _span_enc(enc, mask):
            if mask is None:
                return {k: v.cpu() for k, v in enc.items()}
            return {"input_ids": mask, "attention_mask": mask}

        spans_pos = select_spans(_span_enc(enc_pos, mask_pos), prompt_enc, spec.accumulate)
        spans_neg = select_spans(_span_enc(enc_neg, mask_neg), prompt_enc, spec.accumulate)

        # compute directions via PCA
        directions: dict[int, torch.Tensor] = {}
        explained_variances: dict[int, float] = {}

        num_layers = len(hs_pos)
        logger.debug("Computing directions for %d layers", num_layers)

        for layer_id in range(num_layers):
            # pool over spans
            Hp = pool_over_spans(hs_pos[layer_id], spans_pos)  # [N, H]
            Hn = pool_over_spans(hs_neg[layer_id], spans_neg)  # [N, H]

            samples = _prepare_pca_samples(Hp, Hn, spec.method)  # [2N, H]

            pca = PCA(n_components=1)
            pca.fit(samples.numpy())
            direction = torch.from_numpy(pca.components_[0]).float()  # [H]
            variance = float(pca.explained_variance_ratio_[0])

            direction = _orient_direction(direction, Hp, Hn)
            if not torch.isfinite(direction).all():
                raise ValueError(f"Non-finite direction produced for layer {layer_id}.")

            directions[layer_id] = direction.unsqueeze(0)  # [1, H]
            explained_variances[layer_id] = variance

        logger.debug("Finished fitting contrastive directions")
        meta = artifact_provenance_meta(model, tokenizer) if model is not None else session_meta
        return SteeringVector(
            model_type=model_type,
            directions=directions,
            explained_variances=explained_variances,
            meta=meta,
        )

fit(model, tokenizer, *, data, spec, on_progress=None, session=None)

Extract contrastive direction vectors.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	Model to extract hidden states from, or None to extract through session.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for encoding the contrastive pairs.	required
data	ContrastivePairs	The positive/negative text pairs.	required
spec	VectorTrainSpec	Training configuration (method, accumulate, batch_size).	required
on_progress	Callable[[int, int], None] | None	Optional (completed, total) callback fired as each forward-pass batch finishes. total covers both positive and negative passes.	None
session	SteeringSession | None	A SteeringSession serving hidden-state capture when no live model is available.	None

Returns:

Type	Description
SteeringVector	SteeringVector with one direction per layer.

Source code in steerability/algorithms/state_control/common/estimators/contrastive_direction.py

def fit(
    self,
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase,
    *,
    data: ContrastivePairs,
    spec: VectorTrainSpec,
    on_progress: Callable[[int, int], None] | None = None,
    session: "SteeringSession | None" = None,
) -> SteeringVector:
    """Extract contrastive direction vectors.

    Args:
        model: Model to extract hidden states from, or None to extract through `session`.
        tokenizer: Tokenizer for encoding the contrastive pairs.
        data: The positive/negative text pairs.
        spec: Training configuration (method, accumulate, batch_size).
        on_progress: Optional `(completed, total)` callback fired as each forward-pass batch
            finishes. `total` covers both positive and negative passes.
        session: A `SteeringSession` serving hidden-state capture when no live model is
            available.

    Returns:
        SteeringVector with one direction per layer.
    """
    device = next(model.parameters()).device if model is not None else torch.device("cpu")
    if model is not None:
        model_type = getattr(model.config, "model_type", "unknown")
        session_meta: dict = {}
    else:
        model_type, session_meta = session_artifact_identity(session)

    # render full texts according to prompt_format (shared with inference)
    rendered = render_contrastive(tokenizer, data, spec.prompt_format)

    logger.debug(
        "Tokenizing %d positive and %d negative examples", len(rendered.pos_texts), len(rendered.neg_texts)
    )

    # tokenize
    enc_pos = tokenize_texts(tokenizer, rendered.pos_texts, device, add_special_tokens=rendered.add_special_tokens)
    enc_neg = tokenize_texts(tokenizer, rendered.neg_texts, device, add_special_tokens=rendered.add_special_tokens)

    # tokenize prompts separately if needed for suffix-only
    prompt_enc = None
    if spec.accumulate == "suffix-only" and rendered.prompt_texts is not None:
        prompt_enc = tokenize_texts(
            tokenizer, rendered.prompt_texts, device, add_special_tokens=rendered.add_special_tokens
        )
        prompt_enc = {k: v.cpu() for k, v in prompt_enc.items()}

    # extract hidden states
    logger.debug("Extracting hidden states with batch_size=%d", spec.batch_size)
    n_pos = enc_pos["input_ids"].size(0)
    n_neg = enc_neg["input_ids"].size(0)
    total_batches = math.ceil(n_pos / spec.batch_size) + math.ceil(n_neg / spec.batch_size)
    completed = {"n": 0}

    def _tick() -> None:
        completed["n"] += 1
        if on_progress is not None:
            on_progress(completed["n"], total_batches)

    if on_progress is not None:
        on_progress(0, total_batches)
    hs_pos, mask_pos = capture_hidden(
        enc_pos, model=model, session=session,
        batch_size=spec.batch_size, on_batch=_tick, location=spec.location,
    )
    hs_neg, mask_neg = capture_hidden(
        enc_neg, model=model, session=session,
        batch_size=spec.batch_size, on_batch=_tick, location=spec.location,
    )

    # select spans against the returned layout (spans are mask-derived, so the returned
    # mask keeps them aligned when a remote capture re-packs rows)
    def _span_enc(enc, mask):
        if mask is None:
            return {k: v.cpu() for k, v in enc.items()}
        return {"input_ids": mask, "attention_mask": mask}

    spans_pos = select_spans(_span_enc(enc_pos, mask_pos), prompt_enc, spec.accumulate)
    spans_neg = select_spans(_span_enc(enc_neg, mask_neg), prompt_enc, spec.accumulate)

    # compute directions via PCA
    directions: dict[int, torch.Tensor] = {}
    explained_variances: dict[int, float] = {}

    num_layers = len(hs_pos)
    logger.debug("Computing directions for %d layers", num_layers)

    for layer_id in range(num_layers):
        # pool over spans
        Hp = pool_over_spans(hs_pos[layer_id], spans_pos)  # [N, H]
        Hn = pool_over_spans(hs_neg[layer_id], spans_neg)  # [N, H]

        samples = _prepare_pca_samples(Hp, Hn, spec.method)  # [2N, H]

        pca = PCA(n_components=1)
        pca.fit(samples.numpy())
        direction = torch.from_numpy(pca.components_[0]).float()  # [H]
        variance = float(pca.explained_variance_ratio_[0])

        direction = _orient_direction(direction, Hp, Hn)
        if not torch.isfinite(direction).all():
            raise ValueError(f"Non-finite direction produced for layer {layer_id}.")

        directions[layer_id] = direction.unsqueeze(0)  # [1, H]
        explained_variances[layer_id] = variance

    logger.debug("Finished fitting contrastive directions")
    meta = artifact_provenance_meta(model, tokenizer) if model is not None else session_meta
    return SteeringVector(
        model_type=model_type,
        directions=directions,
        explained_variances=explained_variances,
        meta=meta,
    )

_orient_direction(direction, positive, negative)

Orient direction so the positive class projects above the negative class.

Uses a majority vote over pairs, breaking ties by the sign of the mean projection margin.

Parameters:

Name	Type	Description	Default
direction	Tensor	Direction to orient, shape [H].	required
positive	Tensor	Pooled positive activations, shape [N, H].	required
negative	Tensor	Pooled negative activations, shape [N, H].	required

Returns:

Type	Description
Tensor	The direction, flipped if positives projected below negatives.

Source code in steerability/algorithms/state_control/common/estimators/contrastive_direction.py

def _orient_direction(
    direction: torch.Tensor,
    positive: torch.Tensor,
    negative: torch.Tensor,
) -> torch.Tensor:
    """Orient `direction` so the positive class projects above the negative class.

    Uses a majority vote over pairs, breaking ties by the sign of the mean projection margin.

    Args:
        direction: Direction to orient, shape `[H]`.
        positive: Pooled positive activations, shape `[N, H]`.
        negative: Pooled negative activations, shape `[N, H]`.

    Returns:
        The direction, flipped if positives projected below negatives.
    """
    direction = direction.float()
    positive_projection = positive.float() @ direction
    negative_projection = negative.float() @ direction

    positive_wins = (positive_projection > negative_projection).float().mean()
    if positive_wins < 0.5:
        return -direction
    if positive_wins == 0.5:
        mean_margin = (positive_projection - negative_projection).mean()
        if mean_margin < 0:
            return -direction
    return direction

_prepare_pca_samples(positive, negative, method)

Build the PCA sample matrix from pooled positive/negative activations.

• pca_pairwise: centers each pair (H^+_i, H^-_i) at its midpoint, yielding the two samples ±(H^+_i - H^-_i)/2.
• pca_center: stacks positive and negative activations and centers by their grand mean.

Parameters:

Name	Type	Description	Default
positive	Tensor	Pooled positive activations, shape [N, H].	required
negative	Tensor	Pooled negative activations, shape [N, H].	required
method	PcaMethod	Which sample construction to use.	required

Returns:

Type	Description
Tensor	A float32 sample matrix of shape [2N, H].

Raises:

Type	Description
ValueError	If the shapes disagree, the method is unsupported, or the samples are non-finite.

Source code in steerability/algorithms/state_control/common/estimators/contrastive_direction.py

def _prepare_pca_samples(
    positive: torch.Tensor,
    negative: torch.Tensor,
    method: PcaMethod,
) -> torch.Tensor:
    """Build the PCA sample matrix from pooled positive/negative activations.

    - `pca_pairwise`: centers each pair `(H^+_i, H^-_i)` at its midpoint, yielding the two samples
        `±(H^+_i - H^-_i)/2`.
    - `pca_center`: stacks positive and negative activations and centers by their grand mean.

    Args:
        positive: Pooled positive activations, shape `[N, H]`.
        negative: Pooled negative activations, shape `[N, H]`.
        method: Which sample construction to use.

    Returns:
        A float32 sample matrix of shape `[2N, H]`.

    Raises:
        ValueError: If the shapes disagree, the method is unsupported, or the samples are non-finite.
    """
    if positive.shape != negative.shape:
        raise ValueError(
            "positive and negative pooled activations must have equal shape; "
            f"got {tuple(positive.shape)} and {tuple(negative.shape)}."
        )

    positive = positive.float()
    negative = negative.float()

    if method == "pca_pairwise":
        delta = positive - negative
        samples = torch.cat((0.5 * delta, -0.5 * delta), dim=0)
    elif method == "pca_center":
        stacked = torch.cat((positive, negative), dim=0)
        samples = stacked - stacked.mean(dim=0, keepdim=True)
    else:
        raise ValueError(f"Unknown PCA method: {method!r}.")

    if not torch.isfinite(samples).all():
        raise ValueError("PCA samples contain non-finite values.")
    return samples

mean_difference

Mean difference estimator for CAA steering vectors.

logger = logging.getLogger(__name__) module-attribute

MeanDifferenceEstimator

Bases: BaseEstimator[SteeringVector]

Learns per-layer steering vectors using the mean difference method.

For each layer, computes: v_L = mean(a_L(positive) - a_L(negative))

where activations are extracted at the last non-pad token position of each example (the answer letter in the CAA prompt format).

This differs from ContrastiveDirectionEstimator which uses PCA on the pairwise differences. Mean difference takes the centroid directly, while PCA finds the direction of maximum variance. They converge when difference vectors are nearly collinear.

Examples are rendered via render_for_model according to spec.prompt_format and tokenized with add_special_tokens=False for chat-templated text.

Source code in steerability/algorithms/state_control/common/estimators/mean_difference.py

class MeanDifferenceEstimator(BaseEstimator[SteeringVector]):
    """Learns per-layer steering vectors using the mean difference method.

    For each layer, computes:
        v_L = mean(a_L(positive) - a_L(negative))

    where activations are extracted at the last non-pad token position
    of each example (the answer letter in the CAA prompt format).

    This differs from ContrastiveDirectionEstimator which uses PCA on the
    pairwise differences. Mean difference takes the centroid directly, while
    PCA finds the direction of maximum variance. They converge when difference
    vectors are nearly collinear.

    Examples are rendered via `render_for_model` according to `spec.prompt_format`
    and tokenized with `add_special_tokens=False` for chat-templated text.
    """

    def fit(
        self,
        model: PreTrainedModel | None,
        tokenizer: PreTrainedTokenizerBase,
        *,
        data: ContrastivePairs,
        spec: VectorTrainSpec,
        on_progress: Callable[[int, int], None] | None = None,
        session: "SteeringSession | None" = None,
    ) -> SteeringVector:
        """Extract steering vectors using mean difference.

        Args:
            model: Model to extract hidden states from, or None to extract through `session`.
            tokenizer: Tokenizer for encoding the contrastive pairs.
            data: The positive/negative text pairs.
            spec: Training configuration (method, accumulate, batch_size).
            on_progress: Optional `(completed, total)` callback fired as each forward-pass batch
                finishes. `total` covers both positive and negative passes.
            session: A `SteeringSession` serving hidden-state capture when no live model is
                available.

        Returns:
            SteeringVector with one direction per layer.
        """
        device = next(model.parameters()).device if model is not None else torch.device("cpu")
        if model is not None:
            model_type = getattr(model.config, "model_type", "unknown")
            session_meta: dict = {}
        else:
            model_type, session_meta = session_artifact_identity(session)

        # render full texts according to prompt_format (shared with inference)
        rendered = render_contrastive(tokenizer, data, spec.prompt_format)

        logger.debug(
            "Tokenizing %d positive and %d negative examples", len(rendered.pos_texts), len(rendered.neg_texts)
        )

        # tokenize pairs together to ensure consistent padding and token alignment
        enc_pos, enc_neg = tokenize_pairs(
            tokenizer, rendered.pos_texts, rendered.neg_texts, device, add_special_tokens=rendered.add_special_tokens
        )

        # extract hidden states
        logger.debug("Extracting hidden states with batch_size=%d", spec.batch_size)
        n_pos = enc_pos["input_ids"].size(0)
        n_neg = enc_neg["input_ids"].size(0)
        total_batches = math.ceil(n_pos / spec.batch_size) + math.ceil(n_neg / spec.batch_size)
        completed = {"n": 0}

        def _tick() -> None:
            completed["n"] += 1
            if on_progress is not None:
                on_progress(completed["n"], total_batches)

        if on_progress is not None:
            on_progress(0, total_batches)
        hs_pos, attn_pos = capture_hidden(
            enc_pos, model=model, session=session,
            batch_size=spec.batch_size, on_batch=_tick, location=spec.location,
        )
        hs_neg, attn_neg = capture_hidden(
            enc_neg, model=model, session=session,
            batch_size=spec.batch_size, on_batch=_tick, location=spec.location,
        )

        num_samples = len(rendered.pos_texts)
        num_layers = len(hs_pos)
        logger.debug("Computing mean difference directions for %d layers", num_layers)

        # determine how to aggregate hidden states based on accumulate mode
        directions: dict[int, torch.Tensor] = {}

        for layer_id in range(num_layers):
            hp = hs_pos[layer_id]  # [N, T, H]
            hn = hs_neg[layer_id]  # [N, T, H]

            if spec.accumulate == "last_token":
                # select activation at last non-pad token
                pos_positions = get_last_token_positions(attn_pos, hp.size(1), num_samples)
                neg_positions = get_last_token_positions(attn_neg, hn.size(1), num_samples)
                hp_agg = select_at_positions(hp, pos_positions)  # [N, H]
                hn_agg = select_at_positions(hn, neg_positions)  # [N, H]
            elif spec.accumulate == "all":
                # mean pool over real tokens only; pooling pad positions would bias the direction
                hp_agg = _masked_mean(hp, attn_pos)  # [N, H]
                hn_agg = _masked_mean(hn, attn_neg)  # [N, H]
            else:
                raise ValueError(f"MeanDifferenceEstimator does not support accumulate='{spec.accumulate}'")

            # compute mean difference: v = mean(h_pos - h_neg)
            diffs = hp_agg - hn_agg  # [N, H]
            direction = diffs.mean(dim=0)  # [H]

            directions[layer_id] = direction.unsqueeze(0).to(dtype=torch.float32)  # [1, H]

        logger.debug("Finished fitting mean difference directions")
        meta = artifact_provenance_meta(model, tokenizer) if model is not None else session_meta
        return SteeringVector(
            model_type=model_type,
            directions=directions,
            meta=meta,
        )

fit(model, tokenizer, *, data, spec, on_progress=None, session=None)

Extract steering vectors using mean difference.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	Model to extract hidden states from, or None to extract through session.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for encoding the contrastive pairs.	required
data	ContrastivePairs	The positive/negative text pairs.	required
spec	VectorTrainSpec	Training configuration (method, accumulate, batch_size).	required
on_progress	Callable[[int, int], None] | None	Optional (completed, total) callback fired as each forward-pass batch finishes. total covers both positive and negative passes.	None
session	SteeringSession | None	A SteeringSession serving hidden-state capture when no live model is available.	None

Returns:

Type	Description
SteeringVector	SteeringVector with one direction per layer.

Source code in steerability/algorithms/state_control/common/estimators/mean_difference.py

def fit(
    self,
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase,
    *,
    data: ContrastivePairs,
    spec: VectorTrainSpec,
    on_progress: Callable[[int, int], None] | None = None,
    session: "SteeringSession | None" = None,
) -> SteeringVector:
    """Extract steering vectors using mean difference.

    Args:
        model: Model to extract hidden states from, or None to extract through `session`.
        tokenizer: Tokenizer for encoding the contrastive pairs.
        data: The positive/negative text pairs.
        spec: Training configuration (method, accumulate, batch_size).
        on_progress: Optional `(completed, total)` callback fired as each forward-pass batch
            finishes. `total` covers both positive and negative passes.
        session: A `SteeringSession` serving hidden-state capture when no live model is
            available.

    Returns:
        SteeringVector with one direction per layer.
    """
    device = next(model.parameters()).device if model is not None else torch.device("cpu")
    if model is not None:
        model_type = getattr(model.config, "model_type", "unknown")
        session_meta: dict = {}
    else:
        model_type, session_meta = session_artifact_identity(session)

    # render full texts according to prompt_format (shared with inference)
    rendered = render_contrastive(tokenizer, data, spec.prompt_format)

    logger.debug(
        "Tokenizing %d positive and %d negative examples", len(rendered.pos_texts), len(rendered.neg_texts)
    )

    # tokenize pairs together to ensure consistent padding and token alignment
    enc_pos, enc_neg = tokenize_pairs(
        tokenizer, rendered.pos_texts, rendered.neg_texts, device, add_special_tokens=rendered.add_special_tokens
    )

    # extract hidden states
    logger.debug("Extracting hidden states with batch_size=%d", spec.batch_size)
    n_pos = enc_pos["input_ids"].size(0)
    n_neg = enc_neg["input_ids"].size(0)
    total_batches = math.ceil(n_pos / spec.batch_size) + math.ceil(n_neg / spec.batch_size)
    completed = {"n": 0}

    def _tick() -> None:
        completed["n"] += 1
        if on_progress is not None:
            on_progress(completed["n"], total_batches)

    if on_progress is not None:
        on_progress(0, total_batches)
    hs_pos, attn_pos = capture_hidden(
        enc_pos, model=model, session=session,
        batch_size=spec.batch_size, on_batch=_tick, location=spec.location,
    )
    hs_neg, attn_neg = capture_hidden(
        enc_neg, model=model, session=session,
        batch_size=spec.batch_size, on_batch=_tick, location=spec.location,
    )

    num_samples = len(rendered.pos_texts)
    num_layers = len(hs_pos)
    logger.debug("Computing mean difference directions for %d layers", num_layers)

    # determine how to aggregate hidden states based on accumulate mode
    directions: dict[int, torch.Tensor] = {}

    for layer_id in range(num_layers):
        hp = hs_pos[layer_id]  # [N, T, H]
        hn = hs_neg[layer_id]  # [N, T, H]

        if spec.accumulate == "last_token":
            # select activation at last non-pad token
            pos_positions = get_last_token_positions(attn_pos, hp.size(1), num_samples)
            neg_positions = get_last_token_positions(attn_neg, hn.size(1), num_samples)
            hp_agg = select_at_positions(hp, pos_positions)  # [N, H]
            hn_agg = select_at_positions(hn, neg_positions)  # [N, H]
        elif spec.accumulate == "all":
            # mean pool over real tokens only; pooling pad positions would bias the direction
            hp_agg = _masked_mean(hp, attn_pos)  # [N, H]
            hn_agg = _masked_mean(hn, attn_neg)  # [N, H]
        else:
            raise ValueError(f"MeanDifferenceEstimator does not support accumulate='{spec.accumulate}'")

        # compute mean difference: v = mean(h_pos - h_neg)
        diffs = hp_agg - hn_agg  # [N, H]
        direction = diffs.mean(dim=0)  # [H]

        directions[layer_id] = direction.unsqueeze(0).to(dtype=torch.float32)  # [1, H]

    logger.debug("Finished fitting mean difference directions")
    meta = artifact_provenance_meta(model, tokenizer) if model is not None else session_meta
    return SteeringVector(
        model_type=model_type,
        directions=directions,
        meta=meta,
    )

single_pair

Single-pair estimator for ActAdd steering vectors.

logger = logging.getLogger(__name__) module-attribute

SinglePairEstimator

Bases: BaseEstimator[SteeringVector]

Extracts per-token positional steering vectors from a single prompt pair.

Given one positive prompt and one negative prompt, it computes the per-token activation difference at every layer (or a specified subset of layers) at the layer-input boundary, i.e. the residual stream a layer pre-hook observes, preserving the full positional structure of the contrast.

The prompts are tokenized without special tokens, so row t of the result corresponds to content token t of the (co-padded) pair. The result is a [T, H] direction matrix per layer, where T is the token length of the longer prompt, and the artifact records meta["location"] = "layer_input".

Source code in steerability/algorithms/state_control/common/estimators/single_pair.py

class SinglePairEstimator(BaseEstimator[SteeringVector]):
    """Extracts per-token positional steering vectors from a single prompt pair.

    Given one positive prompt and one negative prompt, it computes the per-token
    activation difference at every layer (or a specified subset of layers) at the
    layer-input boundary, i.e. the residual stream a layer pre-hook observes,
    preserving the full positional structure of the contrast.

    The prompts are tokenized without special tokens, so row `t` of the result
    corresponds to content token `t` of the (co-padded) pair. The result is a
    `[T, H]` direction matrix per layer, where `T` is the token length of the
    longer prompt, and the artifact records `meta["location"] = "layer_input"`.
    """

    def fit(
        self,
        model: PreTrainedModel | None,
        tokenizer: PreTrainedTokenizerBase,
        *,
        positive_prompt: str,
        negative_prompt: str,
        layer_ids: list[int] | None = None,
        session: "SteeringSession | None" = None,
    ) -> SteeringVector:
        """Extract positional steering vector from a single prompt pair.

        Args:
            model: Model to extract hidden states from, or None to extract through `session`.
            tokenizer: Tokenizer for encoding the prompts.
            positive_prompt: Prompt representing the desired direction
                (e.g., "Love", "I talk about weddings constantly").
            negative_prompt: Prompt representing the opposite direction
                (e.g., "Hate", "I do not talk about weddings constantly").
            layer_ids: If provided, only compute directions for these layers.
                If None, compute for all layers.
            session: A `SteeringSession` serving hidden-state capture when no live model is
                available.

        Returns:
            SteeringVector with [T, H] directions per layer.
        """
        device = next(model.parameters()).device if model is not None else torch.device("cpu")
        if model is not None:
            model_type = getattr(model.config, "model_type", "unknown")
            session_meta: dict = {}
        else:
            model_type, session_meta = session_artifact_identity(session)

        logger.debug("Tokenizing prompt pair: positive=%r, negative=%r", positive_prompt, negative_prompt)

        # co-pad the shorter prompt with a space token; pad tokens like EOS produce
        # different activations
        space_ids = tokenizer.encode(" ", add_special_tokens=False)
        original_pad_token_id = tokenizer.pad_token_id
        if space_ids:
            tokenizer.pad_token_id = space_ids[0]
        try:
            enc = tokenizer(
                [positive_prompt, negative_prompt],
                return_tensors="pt",
                padding=True,
                truncation=True,
                add_special_tokens=False,
            )
        finally:
            tokenizer.pad_token_id = original_pad_token_id

        enc = {k: v.to(device) for k, v in enc.items()}

        logger.debug("Running forward pass to extract hidden states")

        hidden, _ = capture_hidden(enc, model=model, session=session, location="layer_input")

        directions: dict[int, torch.Tensor] = {}

        num_layers = len(hidden)
        logger.debug("Computing per-token difference for %d layers", num_layers)

        for layer_idx in range(num_layers):
            if layer_ids is not None and layer_idx not in layer_ids:
                continue

            hs = hidden[layer_idx]  # [2, T, H]
            h_pos = hs[0]  # [T, H]
            h_neg = hs[1]  # [T, H]

            direction = (h_pos - h_neg).cpu().to(dtype=torch.float32)  # [T, H]

            directions[layer_idx] = direction

        logger.debug("Finished fitting single-pair directions with T=%d tokens", direction.size(0))
        meta = artifact_provenance_meta(model, tokenizer) if model is not None else session_meta
        meta["location"] = "layer_input"
        return SteeringVector(
            model_type=model_type,
            directions=directions,
            meta=meta,
        )

fit(model, tokenizer, *, positive_prompt, negative_prompt, layer_ids=None, session=None)

Extract positional steering vector from a single prompt pair.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	Model to extract hidden states from, or None to extract through session.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for encoding the prompts.	required
positive_prompt	str	Prompt representing the desired direction (e.g., "Love", "I talk about weddings constantly").	required
negative_prompt	str	Prompt representing the opposite direction (e.g., "Hate", "I do not talk about weddings constantly").	required
layer_ids	list[int] | None	If provided, only compute directions for these layers. If None, compute for all layers.	None
session	SteeringSession | None	A SteeringSession serving hidden-state capture when no live model is available.	None

Returns:

Type	Description
SteeringVector	SteeringVector with [T, H] directions per layer.

Source code in steerability/algorithms/state_control/common/estimators/single_pair.py

def fit(
    self,
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase,
    *,
    positive_prompt: str,
    negative_prompt: str,
    layer_ids: list[int] | None = None,
    session: "SteeringSession | None" = None,
) -> SteeringVector:
    """Extract positional steering vector from a single prompt pair.

    Args:
        model: Model to extract hidden states from, or None to extract through `session`.
        tokenizer: Tokenizer for encoding the prompts.
        positive_prompt: Prompt representing the desired direction
            (e.g., "Love", "I talk about weddings constantly").
        negative_prompt: Prompt representing the opposite direction
            (e.g., "Hate", "I do not talk about weddings constantly").
        layer_ids: If provided, only compute directions for these layers.
            If None, compute for all layers.
        session: A `SteeringSession` serving hidden-state capture when no live model is
            available.

    Returns:
        SteeringVector with [T, H] directions per layer.
    """
    device = next(model.parameters()).device if model is not None else torch.device("cpu")
    if model is not None:
        model_type = getattr(model.config, "model_type", "unknown")
        session_meta: dict = {}
    else:
        model_type, session_meta = session_artifact_identity(session)

    logger.debug("Tokenizing prompt pair: positive=%r, negative=%r", positive_prompt, negative_prompt)

    # co-pad the shorter prompt with a space token; pad tokens like EOS produce
    # different activations
    space_ids = tokenizer.encode(" ", add_special_tokens=False)
    original_pad_token_id = tokenizer.pad_token_id
    if space_ids:
        tokenizer.pad_token_id = space_ids[0]
    try:
        enc = tokenizer(
            [positive_prompt, negative_prompt],
            return_tensors="pt",
            padding=True,
            truncation=True,
            add_special_tokens=False,
        )
    finally:
        tokenizer.pad_token_id = original_pad_token_id

    enc = {k: v.to(device) for k, v in enc.items()}

    logger.debug("Running forward pass to extract hidden states")

    hidden, _ = capture_hidden(enc, model=model, session=session, location="layer_input")

    directions: dict[int, torch.Tensor] = {}

    num_layers = len(hidden)
    logger.debug("Computing per-token difference for %d layers", num_layers)

    for layer_idx in range(num_layers):
        if layer_ids is not None and layer_idx not in layer_ids:
            continue

        hs = hidden[layer_idx]  # [2, T, H]
        h_pos = hs[0]  # [T, H]
        h_neg = hs[1]  # [T, H]

        direction = (h_pos - h_neg).cpu().to(dtype=torch.float32)  # [T, H]

        directions[layer_idx] = direction

    logger.debug("Finished fitting single-pair directions with T=%d tokens", direction.size(0))
    meta = artifact_provenance_meta(model, tokenizer) if model is not None else session_meta
    meta["location"] = "layer_input"
    return SteeringVector(
        model_type=model_type,
        directions=directions,
        meta=meta,
    )

steering_plane

Steering-plane estimator: per-layer 2D bases for angular steering.

logger = logging.getLogger(__name__) module-attribute

SteeringPlaneEstimator

Bases: BaseEstimator[SteeringVector]

Learns a 2D steering plane per layer (paper sections 4.4-4.5).

The plane is built offline in three steps:

1. Per-layer feature axis computed as the difference in means between positive and negative activations.
2. One global companion axis: the first PCA component across the stacked per-layer feature directions.
3. Per-layer Gram-Schmidt into an orthonormal (b1, b2), with b1 = normalize(d_feat) and b2 = normalize(d_PC0 - (d_PC0·b1) b1).

Returns a SteeringVector with directions[layer] of shape [2, H] (row 0 = feature axis, row 1 = companion axis). The cross-layer PCA explained-variance fraction is stored in explained_variances under the sentinel key -1.

Reference:

• "Angular Steering: Behavior Control via Rotation in Activation Space" Hieu M. Vu, Tan M. Nguyen https://arxiv.org/abs/2510.26243

Source code in steerability/algorithms/state_control/common/estimators/steering_plane.py

class SteeringPlaneEstimator(BaseEstimator[SteeringVector]):
    """Learns a 2D steering plane per layer (paper sections 4.4-4.5).

    The plane is built offline in three steps:

    1. Per-layer feature axis computed as the difference in means between positive and negative
        activations.
    2. One global companion axis: the first PCA component across the stacked per-layer feature
        directions.
    3. Per-layer Gram-Schmidt into an orthonormal `(b1, b2)`, with `b1 = normalize(d_feat)` and
        `b2 = normalize(d_PC0 - (d_PC0·b1) b1)`.

    Returns a `SteeringVector` with `directions[layer]` of shape `[2, H]` (row 0 = feature axis,
    row 1 = companion axis). The cross-layer PCA explained-variance fraction is stored in
    `explained_variances` under the sentinel key `-1`.

    Reference:

    - "Angular Steering: Behavior Control via Rotation in Activation Space"
      Hieu M. Vu, Tan M. Nguyen
      [https://arxiv.org/abs/2510.26243](https://arxiv.org/abs/2510.26243)
    """

    def fit(
        self,
        model: PreTrainedModel | None,
        tokenizer: PreTrainedTokenizerBase,
        *,
        data: ContrastivePairs,
        spec: VectorTrainSpec,
        on_progress: Callable[[int, int], None] | None = None,
        session=None,
    ) -> SteeringVector:
        """Fit the per-layer steering planes.

        Args:
            model: Model to extract hidden states from.
            tokenizer: Tokenizer for encoding the contrastive pairs.
            data: The positive/negative text pairs.
            spec: Training configuration (method, accumulate, batch_size, prompt_format).
            on_progress: Optional `(completed, total)` callback forwarded to the underlying
                mean-difference extraction.

        Returns:
            SteeringVector with a `[2, H]` orthonormal basis per usable layer.

        Raises:
            ValueError: If fewer than two layers have usable feature directions, or if no plane
                survives orthogonalization.
        """
        # step 1: per-layer feature axis (reuse CAA's estimator)
        feature_sv = MeanDifferenceEstimator().fit(
            model, tokenizer, data=data, spec=spec, on_progress=on_progress, session=session
        )

        layer_ids = sorted(feature_sv.directions.keys())
        feature_dirs = {lid: feature_sv.directions[lid].squeeze(0).float() for lid in layer_ids}

        usable = [lid for lid in layer_ids if feature_dirs[lid].norm() >= 1e-8]
        if len(usable) < 2:
            raise ValueError(
                "SteeringPlaneEstimator needs feature directions from at least two layers "
                f"to run cross-layer PCA; got {len(usable)}."
            )

        # step 2: cross-layer PCA -> single companion axis
        feature_matrix = torch.stack([feature_dirs[lid] for lid in usable], dim=0)  # [L, H]
        pca = PCA(n_components=1)
        pca.fit(feature_matrix.numpy())
        pc0 = torch.tensor(pca.components_[0], dtype=torch.float32)
        pc0_variance = float(pca.explained_variance_ratio_[0])

        # step 3: per-layer Gram-Schmidt -> [b1, b2]
        directions: dict[int, torch.Tensor] = {}
        for lid in usable:
            b1 = feature_dirs[lid]
            b1 = b1 / b1.norm()
            b2 = pc0 - (pc0 @ b1) * b1
            if b2.norm() < 1e-8:
                logger.warning("PC0 collinear with feature direction at layer %d; skipping.", lid)
                continue
            b2 = b2 / b2.norm()
            directions[lid] = torch.stack([b1, b2], dim=0)  # [2, H]

        if not directions:
            raise ValueError("No usable steering planes after orthogonalization.")

        logger.debug("Fitted %d steering planes; PC0 explained variance %.4f", len(directions), pc0_variance)
        return SteeringVector(
            model_type=feature_sv.model_type,
            directions=directions,
            explained_variances={-1: pc0_variance},
            meta=dict(feature_sv.meta),
        )

fit(model, tokenizer, *, data, spec, on_progress=None, session=None)

Fit the per-layer steering planes.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	Model to extract hidden states from.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for encoding the contrastive pairs.	required
data	ContrastivePairs	The positive/negative text pairs.	required
spec	VectorTrainSpec	Training configuration (method, accumulate, batch_size, prompt_format).	required
on_progress	Callable[[int, int], None] | None	Optional (completed, total) callback forwarded to the underlying mean-difference extraction.	None

Returns:

Type	Description
SteeringVector	SteeringVector with a [2, H] orthonormal basis per usable layer.

Raises:

Type	Description
ValueError	If fewer than two layers have usable feature directions, or if no plane survives orthogonalization.

Source code in steerability/algorithms/state_control/common/estimators/steering_plane.py

def fit(
    self,
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase,
    *,
    data: ContrastivePairs,
    spec: VectorTrainSpec,
    on_progress: Callable[[int, int], None] | None = None,
    session=None,
) -> SteeringVector:
    """Fit the per-layer steering planes.

    Args:
        model: Model to extract hidden states from.
        tokenizer: Tokenizer for encoding the contrastive pairs.
        data: The positive/negative text pairs.
        spec: Training configuration (method, accumulate, batch_size, prompt_format).
        on_progress: Optional `(completed, total)` callback forwarded to the underlying
            mean-difference extraction.

    Returns:
        SteeringVector with a `[2, H]` orthonormal basis per usable layer.

    Raises:
        ValueError: If fewer than two layers have usable feature directions, or if no plane
            survives orthogonalization.
    """
    # step 1: per-layer feature axis (reuse CAA's estimator)
    feature_sv = MeanDifferenceEstimator().fit(
        model, tokenizer, data=data, spec=spec, on_progress=on_progress, session=session
    )

    layer_ids = sorted(feature_sv.directions.keys())
    feature_dirs = {lid: feature_sv.directions[lid].squeeze(0).float() for lid in layer_ids}

    usable = [lid for lid in layer_ids if feature_dirs[lid].norm() >= 1e-8]
    if len(usable) < 2:
        raise ValueError(
            "SteeringPlaneEstimator needs feature directions from at least two layers "
            f"to run cross-layer PCA; got {len(usable)}."
        )

    # step 2: cross-layer PCA -> single companion axis
    feature_matrix = torch.stack([feature_dirs[lid] for lid in usable], dim=0)  # [L, H]
    pca = PCA(n_components=1)
    pca.fit(feature_matrix.numpy())
    pc0 = torch.tensor(pca.components_[0], dtype=torch.float32)
    pc0_variance = float(pca.explained_variance_ratio_[0])

    # step 3: per-layer Gram-Schmidt -> [b1, b2]
    directions: dict[int, torch.Tensor] = {}
    for lid in usable:
        b1 = feature_dirs[lid]
        b1 = b1 / b1.norm()
        b2 = pc0 - (pc0 @ b1) * b1
        if b2.norm() < 1e-8:
            logger.warning("PC0 collinear with feature direction at layer %d; skipping.", lid)
            continue
        b2 = b2 / b2.norm()
        directions[lid] = torch.stack([b1, b2], dim=0)  # [2, H]

    if not directions:
        raise ValueError("No usable steering planes after orthogonalization.")

    logger.debug("Fitted %d steering planes; PC0 explained variance %.4f", len(directions), pc0_variance)
    return SteeringVector(
        model_type=feature_sv.model_type,
        directions=directions,
        explained_variances={-1: pc0_variance},
        meta=dict(feature_sv.meta),
    )

fit_specs

Configuration for fitting and searching steering artifacts.

Holds the fit-time vocabulary shared by state control components: VectorTrainSpec describes how direction vectors are extracted, ConditionSearchSpec describes how condition points are searched, and the comparator vocabulary (Comparator, CompMode) carries the gate-comparison semantics ("ge" opens when score >= threshold, "le" when score <= threshold). These specs describe how artifacts are produced; the intervention IR in specs.py describes how bound artifacts are applied.

CompMode = Literal['mean', 'last'] module-attribute

Comparator = Literal['ge', 'le'] module-attribute

ConditionSearchSpec dataclass

Configuration for automatic condition point search.

Attributes:

Name	Type	Description
auto_find	bool	If True, run the search during steer(). If False, the user must provide condition_layer_ids and threshold manually.
candidate_layers	Sequence[int] | None	Explicit layer ids to search over. If None, use layer_range.
layer_range	tuple[int, int] | None	0-based (start, end) half-open range of layers to consider. Ignored if candidate_layers is set. Defaults to all layers.
threshold_range	tuple[float, float]	(min, max) for the threshold grid search (half-open, step-exact).
threshold_step	float	Step size for the threshold grid.

Source code in steerability/algorithms/state_control/common/fit_specs.py

@dataclass(frozen=True)
class ConditionSearchSpec:
    """Configuration for automatic condition point search.

    Attributes:
        auto_find: If True, run the search during steer(). If False, the
            user must provide condition_layer_ids and threshold manually.
        candidate_layers: Explicit layer ids to search over. If None, use
            layer_range.
        layer_range: 0-based (start, end) half-open range of layers to consider. Ignored if
            candidate_layers is set. Defaults to all layers.
        threshold_range: (min, max) for the threshold grid search (half-open, step-exact).
        threshold_step: Step size for the threshold grid.
    """

    auto_find: bool = True
    candidate_layers: Sequence[int] | None = None
    layer_range: tuple[int, int] | None = None
    threshold_range: tuple[float, float] = (0.0, 1.0)
    threshold_step: float = 0.01

    def __post_init__(self):
        lo, hi = self.threshold_range
        if lo >= hi:
            raise ValueError(f"threshold_range ({lo}, {hi}): min must be < max.")
        if self.threshold_step <= 0:
            raise ValueError("threshold_step must be > 0.")

auto_find = True class-attribute instance-attribute

candidate_layers = None class-attribute instance-attribute

layer_range = None class-attribute instance-attribute

threshold_range = (0.0, 1.0) class-attribute instance-attribute

threshold_step = 0.01 class-attribute instance-attribute

__init__(auto_find=True, candidate_layers=None, layer_range=None, threshold_range=(0.0, 1.0), threshold_step=0.01)

__post_init__()

Source code in steerability/algorithms/state_control/common/fit_specs.py

def __post_init__(self):
    lo, hi = self.threshold_range
    if lo >= hi:
        raise ValueError(f"threshold_range ({lo}, {hi}): min must be < max.")
    if self.threshold_step <= 0:
        raise ValueError("threshold_step must be > 0.")

VectorTrainSpec dataclass

Configuration for how to train/extract direction vectors.

Attributes:

Name	Type	Description
method	Literal['pca_pairwise', 'pca_center', 'mean_diff']	Extraction algorithm. "pca_pairwise" uses PCA on paired differences of hidden states. "pca_center" uses PCA on all positive/negative hidden states centered by their grand mean (the CAST extraction from the paper). "mean_diff" uses the mean difference of hidden states (CAA method).
accumulate	Literal['all', 'suffix-only', 'last_token']	How to select hidden state spans for aggregation. "all" uses the full sequence. "suffix-only" uses only the portion after the shared prompt. "last_token" uses only the final non-pad token position.
batch_size	int	Batch size for hidden state extraction forward passes.
prompt_format	PromptFormat	How to render contrastive examples into model-ready text (via render_for_model); the rendered string is tokenized with add_special_tokens=False. "chat_completion" renders prompts as user turns and appends positives/negatives as completions (prompt+answer pairs, e.g. CAA); falls back to "raw" when no prompts are provided. "chat_prompt" renders each positive/negative as a standalone user turn (standalone-prompt contrasts, e.g. the CAST condition); matches the inference rendering exactly. "raw" concatenates prompts + text verbatim with no chat template (base-model methods and standalone statements).
location	HiddenStateLocation	Residual-stream boundary each layer key maps to. outputs.hidden_states is a tuple of num_layers + 1 tensors: index 0 is the embedding output (the input to layer 0) and index i is the output of layer i - 1. "layer_output" (default): key l maps to the output of layer l (hidden_states[l + 1]), the boundary hooked by controls that intervene on the layer output. "layer_input": key l maps to the input of layer l, i.e. the output of layer l - 1 (hidden_states[l]), the boundary observed by layer pre-hooks. A vector fit at one boundary is a distinct artifact from one fit at the other, so fit it at the boundary where the consuming control scores or applies it.

Source code in steerability/algorithms/state_control/common/fit_specs.py

@dataclass(frozen=True)
class VectorTrainSpec:
    """Configuration for how to train/extract direction vectors.

    Attributes:
        method: Extraction algorithm.
            "pca_pairwise" uses PCA on paired differences of hidden states.
            "pca_center" uses PCA on all positive/negative hidden states centered
                by their grand mean (the CAST extraction from the paper).
            "mean_diff" uses the mean difference of hidden states (CAA method).
        accumulate: How to select hidden state spans for aggregation.
            "all" uses the full sequence.
            "suffix-only" uses only the portion after the shared prompt.
            "last_token" uses only the final non-pad token position.
        batch_size: Batch size for hidden state extraction forward passes.
        prompt_format: How to render contrastive examples into model-ready text
            (via `render_for_model`); the rendered string is tokenized with
            `add_special_tokens=False`.
            "chat_completion" renders `prompts` as user turns and appends
            positives/negatives as completions (prompt+answer pairs, e.g. CAA);
            falls back to "raw" when no `prompts` are provided.
            "chat_prompt" renders each positive/negative as a standalone user turn
            (standalone-prompt contrasts, e.g. the CAST condition); matches the
            inference rendering exactly.
            "raw" concatenates `prompts` + text verbatim with no chat template
            (base-model methods and standalone statements).
        location: Residual-stream boundary each layer key maps to. `outputs.hidden_states` is a
            tuple of `num_layers + 1` tensors: index 0 is the embedding output (the input to layer
            0) and index `i` is the output of layer `i - 1`.
            "layer_output" (default): key `l` maps to the output of layer `l`
            (`hidden_states[l + 1]`), the boundary hooked by controls that intervene on the layer
            output.
            "layer_input": key `l` maps to the input of layer `l`, i.e. the output of layer `l - 1`
            (`hidden_states[l]`), the boundary observed by layer pre-hooks.
            A vector fit at one boundary is a distinct artifact from one fit at the other, so fit it
            at the boundary where the consuming control scores or applies it.
    """

    method: Literal["pca_pairwise", "pca_center", "mean_diff"] = "pca_pairwise"
    accumulate: Literal["all", "suffix-only", "last_token"] = "all"
    batch_size: int = 8
    prompt_format: PromptFormat = "chat_completion"
    location: HiddenStateLocation = "layer_output"

    def __post_init__(self):
        if self.batch_size < 1:
            raise ValueError("batch_size must be >= 1.")
        if self.prompt_format not in ("raw", "chat_completion", "chat_prompt"):
            raise ValueError(
                f"prompt_format must be one of raw/chat_completion/chat_prompt, got {self.prompt_format!r}."
            )
        if self.location not in ("layer_output", "layer_input"):
            raise ValueError(
                f"location must be 'layer_output' or 'layer_input', got {self.location!r}."
            )

accumulate = 'all' class-attribute instance-attribute

batch_size = 8 class-attribute instance-attribute

location = 'layer_output' class-attribute instance-attribute

method = 'pca_pairwise' class-attribute instance-attribute

prompt_format = 'chat_completion' class-attribute instance-attribute

__init__(method='pca_pairwise', accumulate='all', batch_size=8, prompt_format='chat_completion', location='layer_output')

__post_init__()

Source code in steerability/algorithms/state_control/common/fit_specs.py

def __post_init__(self):
    if self.batch_size < 1:
        raise ValueError("batch_size must be >= 1.")
    if self.prompt_format not in ("raw", "chat_completion", "chat_prompt"):
        raise ValueError(
            f"prompt_format must be one of raw/chat_completion/chat_prompt, got {self.prompt_format!r}."
        )
    if self.location not in ("layer_output", "layer_input"):
        raise ValueError(
            f"location must be 'layer_output' or 'layer_input', got {self.location!r}."
        )

gating

Gating for state control interventions.

A gate is an operation on evidence decided by a rule. Evidence names where hidden states are read (condition layers and pooling) and the Readout that turns pooled hidden states into one value per row; a Rule is a pure decision over the per-layer values; Gate is the single stateful shell that collects evidence during a generation, freezes the decision when the rule reports complete, and answers per-row open/closed queries from behavior hooks. Provenance (probes, contrastive fits, manual weights) produces readouts; it never appears in gate types.

Wire expressibility is a per-component property: a readout or rule with a wire_kind exports its wire form, and Gate.export composes them into the wire gate object. A CallableReadout has no wire form and keeps the intervention in process.

This module also holds the score math shared with selectors/condition_point.py (rank_one_projector, projected_cosine_similarity_tensor), so selector calibration and runtime scoring provably agree.

AffineReadout

Signed linear readout: value = weights[layer] . pooled.

Constructed from a per-layer weight mapping or a SteeringVector; probe-backed gates build one from the probe's weights via gate_from_probe. Returns zeros for layers absent from the weights. The wire form stacks the weight rows aligned with the evidence layer order.

Parameters:

Name	Type	Description	Default
weights	SteeringVector | Mapping[int, Tensor]	Per-layer weight vectors ([H], or [K, H] using row 0), as a SteeringVector or mapping.	required
location	str | None	The boundary the weights were fitted at, or None to skip the boundary check.	None
model_fingerprint	str | None	The fitted model's identity, or None to skip the identity check.	None

Source code in steerability/algorithms/state_control/common/gating.py

class AffineReadout:
    """Signed linear readout: `value = weights[layer] . pooled`.

    Constructed from a per-layer weight mapping or a `SteeringVector`; probe-backed gates build
    one from the probe's weights via `gate_from_probe`. Returns zeros for layers absent from
    the weights. The wire form stacks the weight rows aligned with the evidence layer order.

    Args:
        weights: Per-layer weight vectors (`[H]`, or `[K, H]` using row 0), as a
            `SteeringVector` or mapping.
        location: The boundary the weights were fitted at, or None to skip the boundary check.
        model_fingerprint: The fitted model's identity, or None to skip the identity check.
    """

    wire_kind: ClassVar[str | None] = "affine"

    def __init__(
        self,
        weights: SteeringVector | Mapping[int, torch.Tensor],
        location: str | None = None,
        model_fingerprint: str | None = None,
    ):
        self.weights = {
            lid: _feature_row(direction).detach().reshape(-1).to(torch.float32)
            for lid, direction in _extract_directions(weights, type(self).__name__).items()
        }
        self.location = location
        self.model_fingerprint = model_fingerprint

    @torch.no_grad()
    def __call__(self, pooled: torch.Tensor, layer_id: int) -> torch.Tensor:
        """Return `[B]` values `weights[layer_id] . pooled` (zeros if the layer is absent)."""
        weights = self.weights.get(layer_id)
        if weights is None:
            return torch.zeros(pooled.size(0), dtype=torch.float32)
        return (pooled.to(torch.float32) @ weights.to(pooled.device)).float().cpu()

    def export(self, layer_ids: tuple[int, ...]) -> "WireForm | None":
        """The `affine` wire form: `weights [L, H]` row-aligned with `layer_ids`."""
        from .specs import WireForm

        if any(lid not in self.weights for lid in layer_ids):
            return None
        stacked = torch.stack([self.weights[lid] for lid in layer_ids])
        return WireForm(kind="affine", tensors={"weights": stacked})

location = location instance-attribute

model_fingerprint = model_fingerprint instance-attribute

weights = {lid: _feature_row(direction).detach().reshape(-1).to(torch.float32) for lid, direction in _extract_directions(weights, type(self).__name__).items()} instance-attribute

wire_kind = 'affine' class-attribute

__call__(pooled, layer_id)

Return [B] values weights[layer_id] . pooled (zeros if the layer is absent).

Source code in steerability/algorithms/state_control/common/gating.py

@torch.no_grad()
def __call__(self, pooled: torch.Tensor, layer_id: int) -> torch.Tensor:
    """Return `[B]` values `weights[layer_id] . pooled` (zeros if the layer is absent)."""
    weights = self.weights.get(layer_id)
    if weights is None:
        return torch.zeros(pooled.size(0), dtype=torch.float32)
    return (pooled.to(torch.float32) @ weights.to(pooled.device)).float().cpu()

__init__(weights, location=None, model_fingerprint=None)

Source code in steerability/algorithms/state_control/common/gating.py

def __init__(
    self,
    weights: SteeringVector | Mapping[int, torch.Tensor],
    location: str | None = None,
    model_fingerprint: str | None = None,
):
    self.weights = {
        lid: _feature_row(direction).detach().reshape(-1).to(torch.float32)
        for lid, direction in _extract_directions(weights, type(self).__name__).items()
    }
    self.location = location
    self.model_fingerprint = model_fingerprint

export(layer_ids)

The affine wire form: weights [L, H] row-aligned with layer_ids.

Source code in steerability/algorithms/state_control/common/gating.py

def export(self, layer_ids: tuple[int, ...]) -> "WireForm | None":
    """The `affine` wire form: `weights [L, H]` row-aligned with `layer_ids`."""
    from .specs import WireForm

    if any(lid not in self.weights for lid in layer_ids):
        return None
    stacked = torch.stack([self.weights[lid] for lid in layer_ids])
    return WireForm(kind="affine", tensors={"weights": stacked})

CallableReadout

Escape hatch wrapping any (pooled [B, H], layer_id) -> Tensor [B] callable.

Has no wire form, so a gate built on it never lowers; support verdicts attribute the in-process-only status to the readout.

Parameters:

Name	Type	Description	Default
fn	Callable[[Tensor, int], Tensor]	The readout callable.	required
location	str | None	The boundary the callable expects features at, or None to skip the check.	None
model_fingerprint	str | None	The fitted model's identity, or None to skip the check.	None

Source code in steerability/algorithms/state_control/common/gating.py

class CallableReadout:
    """Escape hatch wrapping any `(pooled [B, H], layer_id) -> Tensor [B]` callable.

    Has no wire form, so a gate built on it never lowers; support verdicts attribute the
    in-process-only status to the readout.

    Args:
        fn: The readout callable.
        location: The boundary the callable expects features at, or None to skip the check.
        model_fingerprint: The fitted model's identity, or None to skip the check.
    """

    wire_kind: ClassVar[str | None] = None

    def __init__(
        self,
        fn: Callable[[torch.Tensor, int], torch.Tensor],
        location: str | None = None,
        model_fingerprint: str | None = None,
    ):
        self.fn = fn
        self.location = location
        self.model_fingerprint = model_fingerprint

    def __call__(self, pooled: torch.Tensor, layer_id: int) -> torch.Tensor:
        return self.fn(pooled, layer_id)

    def export(self, layer_ids: tuple[int, ...]) -> "WireForm | None":
        """None; an arbitrary callable has no wire form."""
        return None

fn = fn instance-attribute

location = location instance-attribute

model_fingerprint = model_fingerprint instance-attribute

wire_kind = None class-attribute

__call__(pooled, layer_id)

Source code in steerability/algorithms/state_control/common/gating.py

def __call__(self, pooled: torch.Tensor, layer_id: int) -> torch.Tensor:
    return self.fn(pooled, layer_id)

__init__(fn, location=None, model_fingerprint=None)

Source code in steerability/algorithms/state_control/common/gating.py

def __init__(
    self,
    fn: Callable[[torch.Tensor, int], torch.Tensor],
    location: str | None = None,
    model_fingerprint: str | None = None,
):
    self.fn = fn
    self.location = location
    self.model_fingerprint = model_fingerprint

export(layer_ids)

None; an arbitrary callable has no wire form.

Source code in steerability/algorithms/state_control/common/gating.py

def export(self, layer_ids: tuple[int, ...]) -> "WireForm | None":
    """None; an arbitrary callable has no wire form."""
    return None

CosineReadout

Signed cosine similarity between the pooled state and a per-layer direction.

[K, H] artifacts use row 0. Returns zeros for layers absent from the directions. The score is signed: for a mean-difference direction, positives score high, negatives score low, and unrelated content scores near zero, so a "ge" threshold fails closed on out-of-distribution inputs. Prefer this readout for topic or domain gates whose calibration negatives cannot cover the deployment input space; signed scores can be negative, so pair it with a threshold range that admits negative values, e.g. (-1.0, 1.0).

Parameters:

Name	Type	Description	Default
directions	SteeringVector | Mapping[int, Tensor]	Per-layer directions, a SteeringVector or Mapping[int, torch.Tensor].	required
location	str | None	The boundary the directions were fitted at, or None to skip the check.	None
model_fingerprint	str | None	The fitted model's identity, or None to skip the check.	None

Reference:

• "Steering Llama 2 via Contrastive Activation Addition" Nina Panickssery, Nick Gabrieli, Julian Schulz, Meg Tong, Evan Hubinger, Alexander Matt Turner https://arxiv.org/abs/2312.06681

Source code in steerability/algorithms/state_control/common/gating.py

class CosineReadout:
    """Signed cosine similarity between the pooled state and a per-layer direction.

    `[K, H]` artifacts use row 0. Returns zeros for layers absent from the directions. The
    score is signed: for a mean-difference direction, positives score high, negatives score
    low, and unrelated content scores near zero, so a `"ge"` threshold fails closed on
    out-of-distribution inputs. Prefer this readout for topic or domain gates whose calibration
    negatives cannot cover the deployment input space; signed scores can be negative, so pair
    it with a threshold range that admits negative values, e.g. `(-1.0, 1.0)`.

    Args:
        directions: Per-layer directions, a `SteeringVector` or `Mapping[int, torch.Tensor]`.
        location: The boundary the directions were fitted at, or None to skip the check.
        model_fingerprint: The fitted model's identity, or None to skip the check.

    Reference:

    - "Steering Llama 2 via Contrastive Activation Addition"
      Nina Panickssery, Nick Gabrieli, Julian Schulz, Meg Tong, Evan Hubinger, Alexander Matt Turner
      [https://arxiv.org/abs/2312.06681](https://arxiv.org/abs/2312.06681)
    """

    wire_kind: ClassVar[str | None] = "cosine"

    def __init__(
        self,
        directions: SteeringVector | Mapping[int, torch.Tensor],
        location: str | None = None,
        model_fingerprint: str | None = None,
    ):
        self.directions = _extract_directions(directions, type(self).__name__)
        self.location = location
        self.model_fingerprint = model_fingerprint

    @torch.no_grad()
    def __call__(self, pooled: torch.Tensor, layer_id: int) -> torch.Tensor:
        """Return `[B]` signed cosine values for `layer_id` (zeros if the layer is absent)."""
        direction = self.directions.get(layer_id)
        if direction is None:
            return torch.zeros(pooled.size(0), dtype=torch.float32)
        direction = _feature_row(direction).to(dtype=pooled.dtype, device=pooled.device)
        return F.cosine_similarity(pooled, direction.unsqueeze(0), dim=-1).float().cpu()

    def export(self, layer_ids: tuple[int, ...]) -> "WireForm | None":
        """The `cosine` wire form: `directions [L, H]` row-aligned with `layer_ids`."""
        from .specs import WireForm

        if any(lid not in self.directions for lid in layer_ids):
            return None
        stacked = torch.stack(
            [_feature_row(self.directions[lid]).reshape(-1).to(torch.float32) for lid in layer_ids]
        )
        return WireForm(kind="cosine", tensors={"directions": stacked})

directions = _extract_directions(directions, type(self).__name__) instance-attribute

location = location instance-attribute

model_fingerprint = model_fingerprint instance-attribute

wire_kind = 'cosine' class-attribute

__call__(pooled, layer_id)

Return [B] signed cosine values for layer_id (zeros if the layer is absent).

Source code in steerability/algorithms/state_control/common/gating.py

@torch.no_grad()
def __call__(self, pooled: torch.Tensor, layer_id: int) -> torch.Tensor:
    """Return `[B]` signed cosine values for `layer_id` (zeros if the layer is absent)."""
    direction = self.directions.get(layer_id)
    if direction is None:
        return torch.zeros(pooled.size(0), dtype=torch.float32)
    direction = _feature_row(direction).to(dtype=pooled.dtype, device=pooled.device)
    return F.cosine_similarity(pooled, direction.unsqueeze(0), dim=-1).float().cpu()

__init__(directions, location=None, model_fingerprint=None)

Source code in steerability/algorithms/state_control/common/gating.py

def __init__(
    self,
    directions: SteeringVector | Mapping[int, torch.Tensor],
    location: str | None = None,
    model_fingerprint: str | None = None,
):
    self.directions = _extract_directions(directions, type(self).__name__)
    self.location = location
    self.model_fingerprint = model_fingerprint

export(layer_ids)

The cosine wire form: directions [L, H] row-aligned with layer_ids.

Source code in steerability/algorithms/state_control/common/gating.py

def export(self, layer_ids: tuple[int, ...]) -> "WireForm | None":
    """The `cosine` wire form: `directions [L, H]` row-aligned with `layer_ids`."""
    from .specs import WireForm

    if any(lid not in self.directions for lid in layer_ids):
        return None
    stacked = torch.stack(
        [_feature_row(self.directions[lid]).reshape(-1).to(torch.float32) for lid in layer_ids]
    )
    return WireForm(kind="cosine", tensors={"directions": stacked})

Evidence dataclass

Where gate evidence is read and how hidden states become values.

Attributes:

Name	Type	Description
layer_ids	tuple[int, ...]	Condition layers (0-based decoder-layer indices at the intervention's boundary).
readout	Readout	Per-layer map from pooled hidden states to per-row values.
pooling	Literal['mean', 'last']	Token aggregation over real (non-pad) prompt positions: "mean" pools, "last" selects the last real token per row.

Source code in steerability/algorithms/state_control/common/gating.py

@dataclass(frozen=True)
class Evidence:
    """Where gate evidence is read and how hidden states become values.

    Attributes:
        layer_ids: Condition layers (0-based decoder-layer indices at the intervention's
            boundary).
        readout: Per-layer map from pooled hidden states to per-row values.
        pooling: Token aggregation over real (non-pad) prompt positions: `"mean"` pools,
            `"last"` selects the last real token per row.
    """

    layer_ids: tuple[int, ...]
    readout: Readout
    pooling: Literal["mean", "last"] = "mean"

    def __post_init__(self):
        object.__setattr__(self, "layer_ids", tuple(int(lid) for lid in self.layer_ids))
        if not self.layer_ids:
            raise ValueError("Evidence requires at least one condition layer.")
        if self.pooling not in ("mean", "last"):
            raise ValueError(f"pooling must be 'mean' or 'last'; got {self.pooling!r}.")

layer_ids instance-attribute

pooling = 'mean' class-attribute instance-attribute

readout instance-attribute

__init__(layer_ids, readout, pooling='mean')

__post_init__()

Source code in steerability/algorithms/state_control/common/gating.py

def __post_init__(self):
    object.__setattr__(self, "layer_ids", tuple(int(lid) for lid in self.layer_ids))
    if not self.layer_ids:
        raise ValueError("Evidence requires at least one condition layer.")
    if self.pooling not in ("mean", "last"):
        raise ValueError(f"pooling must be 'mean' or 'last'; got {self.pooling!r}.")

Gate

The stateful shell around evidence and a rule; holds one decision per logical batch row.

Lifecycle per generation call:

1. `reset(num_rows)` clears state from the previous generation and sizes the gate to
   the logical batch (number of prompts, not the beam-expanded batch). Reset is
   idempotent: re-resetting an already-reset gate to the same size leaves it in the
   same cleared state, which shared-instance composition (one gate read by several
   interventions) relies on.
2. `update(values, key=layer_id)` records one condition layer's per-row values as
   evidence arrives from condition hooks. When the rule first reports complete
   (typically once every evidence layer has reported on the prefill pass), the
   decision is computed once, frozen for the generation, and the stored values are
   dropped after copying into the diagnostics snapshot.
3. `open_rows()` is queried by behavior hooks; all rows are closed before evidence.

is_ready() reports True once the decision is frozen; the runtime stops condition scoring then, so the prompt is scored once and the decision holds. A rule whose is_complete never returns True re-scores every pass (in-process only; wire gates are prompt-decided by construction).

Parameters:

Name	Type	Description	Default
evidence	Evidence	Where evidence is read and how hidden states become values.	required
rule	Rule	The decision over the collected values.	required

Source code in steerability/algorithms/state_control/common/gating.py

class Gate:
    """The stateful shell around evidence and a rule; holds one decision per logical batch row.

    Lifecycle per generation call:

        1. `reset(num_rows)` clears state from the previous generation and sizes the gate to
           the logical batch (number of prompts, not the beam-expanded batch). Reset is
           idempotent: re-resetting an already-reset gate to the same size leaves it in the
           same cleared state, which shared-instance composition (one gate read by several
           interventions) relies on.
        2. `update(values, key=layer_id)` records one condition layer's per-row values as
           evidence arrives from condition hooks. When the rule first reports complete
           (typically once every evidence layer has reported on the prefill pass), the
           decision is computed once, frozen for the generation, and the stored values are
           dropped after copying into the diagnostics snapshot.
        3. `open_rows()` is queried by behavior hooks; all rows are closed before evidence.

    `is_ready()` reports True once the decision is frozen; the runtime stops condition scoring
    then, so the prompt is scored once and the decision holds. A rule whose `is_complete`
    never returns True re-scores every pass (in-process only; wire gates are prompt-decided by
    construction).

    Args:
        evidence: Where evidence is read and how hidden states become values.
        rule: The decision over the collected values.
    """

    def __init__(self, evidence: Evidence, rule: Rule):
        self.evidence = evidence
        self.rule = rule
        self.num_rows: int = 1
        self._values: dict[int, torch.Tensor] = {}
        self._decision: torch.BoolTensor | None = None
        self._snapshot: dict[int, torch.Tensor] = {}

    def reset(self, num_rows: int = 1) -> None:
        """Clear all per-generation state and size the gate to `num_rows` logical rows."""
        if num_rows < 1:
            raise ValueError(f"num_rows must be >= 1; got {num_rows}.")
        self.num_rows = int(num_rows)
        self._values.clear()
        self._decision = None
        self._snapshot = {}

    def _coerce_values(self, values: torch.Tensor | float) -> torch.Tensor:
        """Normalize `values` to a float32 `[num_rows]` CPU tensor, enforcing the row contract."""
        if isinstance(values, (int, float)):
            if self.num_rows != 1:
                raise ValueError(
                    f"Gate has {self.num_rows} rows but received a scalar value; readouts "
                    f"must return per-row values ([num_rows]) for batched generation."
                )
            return torch.tensor([float(values)], dtype=torch.float32)
        t = torch.as_tensor(values, dtype=torch.float32).reshape(-1).cpu()
        if t.numel() != self.num_rows:
            raise ValueError(
                f"Gate has {self.num_rows} rows but received {t.numel()} values."
            )
        return t

    def update(self, values: torch.Tensor | float, *, key: int) -> None:
        """Record one condition layer's per-row values; freeze the decision on completion.

        Updates after the freeze are ignored.

        Args:
            values: Per-row values of shape `[num_rows]` (a bare float is accepted only when
                `num_rows == 1`).
            key: The condition layer id the values belong to.
        """
        if self._decision is not None:
            return
        self._values[int(key)] = self._coerce_values(values)
        expected = frozenset(self.evidence.layer_ids)
        if self.rule.is_complete(frozenset(self._values), expected):
            decision = self.rule.decide(self._values, self.num_rows)
            self._decision = torch.as_tensor(decision, dtype=torch.bool).reshape(-1)
            self._snapshot = {lid: rows.clone() for lid, rows in self._values.items()}
            self._values.clear()

    def open_rows(self) -> torch.BoolTensor:
        """Per-row decision of shape `[num_rows]`.

        Returns the frozen decision once the rule has reported complete; before that, the rule
        is evaluated live over the values collected so far (all-closed before any evidence,
        since both shipped rules decide all-closed on empty values).
        """
        if self._decision is not None:
            return self._decision
        decision = self.rule.decide(dict(self._values), self.num_rows)
        return torch.as_tensor(decision, dtype=torch.bool).reshape(-1)

    def is_open(self) -> bool:
        """Scalar convenience: True if any row is open (exact for `num_rows == 1`)."""
        return bool(self.open_rows().any())

    def is_ready(self) -> bool:
        """True once the decision is frozen."""
        return self._decision is not None

    def evidence_values(self) -> dict[int, torch.Tensor]:
        """Per-layer value tensors (`[num_rows]` each) behind the frozen decision.

        Empty before the freeze; retained for diagnostics after the stored values are dropped.
        """
        return dict(self._snapshot)

    def wire_kinds(self) -> tuple[frozenset[str], frozenset[str]] | None:
        """The `({readout kind}, {rule kind})` pair, or None when either has no wire form."""
        readout_kind = type(self.evidence.readout).wire_kind
        rule_kind = type(self.rule).wire_kind
        if readout_kind is None or rule_kind is None:
            return None
        return frozenset({readout_kind}), frozenset({rule_kind})

    def export(self, register: Callable[[Mapping[str, torch.Tensor]], str]) -> dict | None:
        """The wire gate object, or None when the configuration has no wire form.

        The readout's tensors are content-addressed through `register`; rule params inline.
        The caller maps `layers` to wire indices before export, so the layer values here are
        the caller's.

        Args:
            register: Callable mapping a tensor payload to its content-addressed artifact id.

        Returns:
            The wire gate dict (`layers`, `pooling`, `readout`, `rule`), with `layers` left as
            the evidence layers for the caller to remap, or None.
        """
        readout_form = self.evidence.readout.export(self.evidence.layer_ids)
        rule_form = self.rule.export()
        if readout_form is None or rule_form is None:
            return None
        readout_wire: dict = {"kind": readout_form.kind, **readout_form.params}
        if readout_form.tensors:
            readout_wire["artifact"] = register(readout_form.tensors)
        return {
            "layers": [int(lid) for lid in self.evidence.layer_ids],
            "pooling": self.evidence.pooling,
            "readout": readout_wire,
            "rule": {"kind": rule_form.kind, **rule_form.params},
        }

    def to_config(self) -> tuple[dict, dict[str, torch.Tensor] | None]:
        """The gate's serialized form: `(params, readout tensors)`.

        Params carry the evidence layers, pooling, the readout's kind, boundary, and recorded
        model fingerprint, and the rule's kind and parameters; the readout's per-layer tensors
        are returned stacked row-aligned with the evidence layers, using the readout's wire
        names.

        Raises:
            ValueError: If the readout or rule has no wire form (e.g. `CallableReadout`), or
                the readout lacks a tensor for an evidence layer.
        """
        readout = self.evidence.readout
        readout_form = readout.export(self.evidence.layer_ids)
        rule_form = self.rule.export()
        if readout_form is None or rule_form is None:
            offender = type(readout).__name__ if readout_form is None else type(self.rule).__name__
            raise ValueError(
                f"Gate over {offender} has no serialized form; gates serialize only when "
                "their readout and rule have wire kinds."
            )
        params = {
            "layers": [int(lid) for lid in self.evidence.layer_ids],
            "pooling": self.evidence.pooling,
            "readout": {
                "kind": readout_form.kind,
                "location": getattr(readout, "location", None),
                "model_fingerprint": getattr(readout, "model_fingerprint", None),
            },
            "rule": {"kind": rule_form.kind, **rule_form.params},
        }
        return params, dict(readout_form.tensors) or None

    @classmethod
    def from_config(cls, params: dict, *, readout_tensors: Mapping[str, torch.Tensor] | None = None) -> "Gate":
        """Rebuild a gate from its serialized form.

        Args:
            params: The serialized params (`layers`, `pooling`, `readout`, `rule`).
            readout_tensors: The readout's stacked tensors, row-aligned with `params["layers"]`.

        Returns:
            The assembled `Gate`.

        Raises:
            ValueError: If the readout or rule kind is unknown, or the stacked tensor rows do
                not match the layer count.
        """
        layer_ids = [int(lid) for lid in params["layers"]]
        readout_params = params["readout"]
        rule_params = dict(params["rule"])

        readout_classes = {"affine": AffineReadout, "cosine": CosineReadout,
                           "projected_cosine": ProjectedCosineReadout}
        readout_cls = readout_classes.get(readout_params["kind"])
        if readout_cls is None:
            raise ValueError(f"Unknown gate readout kind {readout_params['kind']!r}.")
        tensor_name = "weights" if readout_params["kind"] == "affine" else "directions"
        if readout_tensors is None or tensor_name not in readout_tensors:
            raise ValueError(f"Gate readout tensors missing {tensor_name!r}.")
        stacked = readout_tensors[tensor_name]
        if stacked.size(0) != len(layer_ids):
            raise ValueError(
                f"Gate readout tensor has {stacked.size(0)} rows for {len(layer_ids)} layers."
            )
        per_layer = {lid: stacked[i] for i, lid in enumerate(layer_ids)}
        readout = readout_cls(
            per_layer,
            location=readout_params.get("location"),
            model_fingerprint=readout_params.get("model_fingerprint"),
        )

        rule_kind = rule_params.pop("kind")
        if rule_kind == "sum_threshold":
            rule: Rule = SumThreshold(**rule_params)
        elif rule_kind == "per_key_threshold":
            rule = PerKeyThreshold(**rule_params)
        else:
            raise ValueError(f"Unknown gate rule kind {rule_kind!r}.")

        return cls(Evidence(tuple(layer_ids), readout, pooling=params.get("pooling", "mean")), rule)

_decision = None instance-attribute

_snapshot = {} instance-attribute

_values = {} instance-attribute

evidence = evidence instance-attribute

num_rows = 1 instance-attribute

rule = rule instance-attribute

__init__(evidence, rule)

Source code in steerability/algorithms/state_control/common/gating.py

def __init__(self, evidence: Evidence, rule: Rule):
    self.evidence = evidence
    self.rule = rule
    self.num_rows: int = 1
    self._values: dict[int, torch.Tensor] = {}
    self._decision: torch.BoolTensor | None = None
    self._snapshot: dict[int, torch.Tensor] = {}

_coerce_values(values)

Normalize values to a float32 [num_rows] CPU tensor, enforcing the row contract.

Source code in steerability/algorithms/state_control/common/gating.py

def _coerce_values(self, values: torch.Tensor | float) -> torch.Tensor:
    """Normalize `values` to a float32 `[num_rows]` CPU tensor, enforcing the row contract."""
    if isinstance(values, (int, float)):
        if self.num_rows != 1:
            raise ValueError(
                f"Gate has {self.num_rows} rows but received a scalar value; readouts "
                f"must return per-row values ([num_rows]) for batched generation."
            )
        return torch.tensor([float(values)], dtype=torch.float32)
    t = torch.as_tensor(values, dtype=torch.float32).reshape(-1).cpu()
    if t.numel() != self.num_rows:
        raise ValueError(
            f"Gate has {self.num_rows} rows but received {t.numel()} values."
        )
    return t

evidence_values()

Per-layer value tensors ([num_rows] each) behind the frozen decision.

Empty before the freeze; retained for diagnostics after the stored values are dropped.

Source code in steerability/algorithms/state_control/common/gating.py

def evidence_values(self) -> dict[int, torch.Tensor]:
    """Per-layer value tensors (`[num_rows]` each) behind the frozen decision.

    Empty before the freeze; retained for diagnostics after the stored values are dropped.
    """
    return dict(self._snapshot)

export(register)

The wire gate object, or None when the configuration has no wire form.

The readout's tensors are content-addressed through register; rule params inline. The caller maps layers to wire indices before export, so the layer values here are the caller's.

Parameters:

Name	Type	Description	Default
register	Callable[[Mapping[str, Tensor]], str]	Callable mapping a tensor payload to its content-addressed artifact id.	required

Returns:

Type	Description
dict | None	The wire gate dict (layers, pooling, readout, rule), with layers left as
dict | None	the evidence layers for the caller to remap, or None.

Source code in steerability/algorithms/state_control/common/gating.py

def export(self, register: Callable[[Mapping[str, torch.Tensor]], str]) -> dict | None:
    """The wire gate object, or None when the configuration has no wire form.

    The readout's tensors are content-addressed through `register`; rule params inline.
    The caller maps `layers` to wire indices before export, so the layer values here are
    the caller's.

    Args:
        register: Callable mapping a tensor payload to its content-addressed artifact id.

    Returns:
        The wire gate dict (`layers`, `pooling`, `readout`, `rule`), with `layers` left as
        the evidence layers for the caller to remap, or None.
    """
    readout_form = self.evidence.readout.export(self.evidence.layer_ids)
    rule_form = self.rule.export()
    if readout_form is None or rule_form is None:
        return None
    readout_wire: dict = {"kind": readout_form.kind, **readout_form.params}
    if readout_form.tensors:
        readout_wire["artifact"] = register(readout_form.tensors)
    return {
        "layers": [int(lid) for lid in self.evidence.layer_ids],
        "pooling": self.evidence.pooling,
        "readout": readout_wire,
        "rule": {"kind": rule_form.kind, **rule_form.params},
    }

from_config(params, *, readout_tensors=None) classmethod

Rebuild a gate from its serialized form.

Parameters:

Name	Type	Description	Default
params	dict	The serialized params (layers, pooling, readout, rule).	required
readout_tensors	Mapping[str, Tensor] | None	The readout's stacked tensors, row-aligned with params["layers"].	None

Returns:

Type	Description
'Gate'	The assembled Gate.

Raises:

Type	Description
ValueError	If the readout or rule kind is unknown, or the stacked tensor rows do not match the layer count.

Source code in steerability/algorithms/state_control/common/gating.py

@classmethod
def from_config(cls, params: dict, *, readout_tensors: Mapping[str, torch.Tensor] | None = None) -> "Gate":
    """Rebuild a gate from its serialized form.

    Args:
        params: The serialized params (`layers`, `pooling`, `readout`, `rule`).
        readout_tensors: The readout's stacked tensors, row-aligned with `params["layers"]`.

    Returns:
        The assembled `Gate`.

    Raises:
        ValueError: If the readout or rule kind is unknown, or the stacked tensor rows do
            not match the layer count.
    """
    layer_ids = [int(lid) for lid in params["layers"]]
    readout_params = params["readout"]
    rule_params = dict(params["rule"])

    readout_classes = {"affine": AffineReadout, "cosine": CosineReadout,
                       "projected_cosine": ProjectedCosineReadout}
    readout_cls = readout_classes.get(readout_params["kind"])
    if readout_cls is None:
        raise ValueError(f"Unknown gate readout kind {readout_params['kind']!r}.")
    tensor_name = "weights" if readout_params["kind"] == "affine" else "directions"
    if readout_tensors is None or tensor_name not in readout_tensors:
        raise ValueError(f"Gate readout tensors missing {tensor_name!r}.")
    stacked = readout_tensors[tensor_name]
    if stacked.size(0) != len(layer_ids):
        raise ValueError(
            f"Gate readout tensor has {stacked.size(0)} rows for {len(layer_ids)} layers."
        )
    per_layer = {lid: stacked[i] for i, lid in enumerate(layer_ids)}
    readout = readout_cls(
        per_layer,
        location=readout_params.get("location"),
        model_fingerprint=readout_params.get("model_fingerprint"),
    )

    rule_kind = rule_params.pop("kind")
    if rule_kind == "sum_threshold":
        rule: Rule = SumThreshold(**rule_params)
    elif rule_kind == "per_key_threshold":
        rule = PerKeyThreshold(**rule_params)
    else:
        raise ValueError(f"Unknown gate rule kind {rule_kind!r}.")

    return cls(Evidence(tuple(layer_ids), readout, pooling=params.get("pooling", "mean")), rule)

is_open()

Scalar convenience: True if any row is open (exact for num_rows == 1).

Source code in steerability/algorithms/state_control/common/gating.py

def is_open(self) -> bool:
    """Scalar convenience: True if any row is open (exact for `num_rows == 1`)."""
    return bool(self.open_rows().any())

is_ready()

True once the decision is frozen.

Source code in steerability/algorithms/state_control/common/gating.py

def is_ready(self) -> bool:
    """True once the decision is frozen."""
    return self._decision is not None

open_rows()

Per-row decision of shape [num_rows].

Returns the frozen decision once the rule has reported complete; before that, the rule is evaluated live over the values collected so far (all-closed before any evidence, since both shipped rules decide all-closed on empty values).

Source code in steerability/algorithms/state_control/common/gating.py

def open_rows(self) -> torch.BoolTensor:
    """Per-row decision of shape `[num_rows]`.

    Returns the frozen decision once the rule has reported complete; before that, the rule
    is evaluated live over the values collected so far (all-closed before any evidence,
    since both shipped rules decide all-closed on empty values).
    """
    if self._decision is not None:
        return self._decision
    decision = self.rule.decide(dict(self._values), self.num_rows)
    return torch.as_tensor(decision, dtype=torch.bool).reshape(-1)

reset(num_rows=1)

Clear all per-generation state and size the gate to num_rows logical rows.

Source code in steerability/algorithms/state_control/common/gating.py

def reset(self, num_rows: int = 1) -> None:
    """Clear all per-generation state and size the gate to `num_rows` logical rows."""
    if num_rows < 1:
        raise ValueError(f"num_rows must be >= 1; got {num_rows}.")
    self.num_rows = int(num_rows)
    self._values.clear()
    self._decision = None
    self._snapshot = {}

to_config()

The gate's serialized form: (params, readout tensors).

Params carry the evidence layers, pooling, the readout's kind, boundary, and recorded model fingerprint, and the rule's kind and parameters; the readout's per-layer tensors are returned stacked row-aligned with the evidence layers, using the readout's wire names.

Raises:

Type	Description
ValueError	If the readout or rule has no wire form (e.g. CallableReadout), or the readout lacks a tensor for an evidence layer.

Source code in steerability/algorithms/state_control/common/gating.py

def to_config(self) -> tuple[dict, dict[str, torch.Tensor] | None]:
    """The gate's serialized form: `(params, readout tensors)`.

    Params carry the evidence layers, pooling, the readout's kind, boundary, and recorded
    model fingerprint, and the rule's kind and parameters; the readout's per-layer tensors
    are returned stacked row-aligned with the evidence layers, using the readout's wire
    names.

    Raises:
        ValueError: If the readout or rule has no wire form (e.g. `CallableReadout`), or
            the readout lacks a tensor for an evidence layer.
    """
    readout = self.evidence.readout
    readout_form = readout.export(self.evidence.layer_ids)
    rule_form = self.rule.export()
    if readout_form is None or rule_form is None:
        offender = type(readout).__name__ if readout_form is None else type(self.rule).__name__
        raise ValueError(
            f"Gate over {offender} has no serialized form; gates serialize only when "
            "their readout and rule have wire kinds."
        )
    params = {
        "layers": [int(lid) for lid in self.evidence.layer_ids],
        "pooling": self.evidence.pooling,
        "readout": {
            "kind": readout_form.kind,
            "location": getattr(readout, "location", None),
            "model_fingerprint": getattr(readout, "model_fingerprint", None),
        },
        "rule": {"kind": rule_form.kind, **rule_form.params},
    }
    return params, dict(readout_form.tensors) or None

update(values, *, key)

Record one condition layer's per-row values; freeze the decision on completion.

Updates after the freeze are ignored.

Parameters:

Name	Type	Description	Default
values	Tensor | float	Per-row values of shape [num_rows] (a bare float is accepted only when num_rows == 1).	required
key	int	The condition layer id the values belong to.	required

Source code in steerability/algorithms/state_control/common/gating.py

def update(self, values: torch.Tensor | float, *, key: int) -> None:
    """Record one condition layer's per-row values; freeze the decision on completion.

    Updates after the freeze are ignored.

    Args:
        values: Per-row values of shape `[num_rows]` (a bare float is accepted only when
            `num_rows == 1`).
        key: The condition layer id the values belong to.
    """
    if self._decision is not None:
        return
    self._values[int(key)] = self._coerce_values(values)
    expected = frozenset(self.evidence.layer_ids)
    if self.rule.is_complete(frozenset(self._values), expected):
        decision = self.rule.decide(self._values, self.num_rows)
        self._decision = torch.as_tensor(decision, dtype=torch.bool).reshape(-1)
        self._snapshot = {lid: rows.clone() for lid, rows in self._values.items()}
        self._values.clear()

wire_kinds()

The ({readout kind}, {rule kind}) pair, or None when either has no wire form.

Source code in steerability/algorithms/state_control/common/gating.py

def wire_kinds(self) -> tuple[frozenset[str], frozenset[str]] | None:
    """The `({readout kind}, {rule kind})` pair, or None when either has no wire form."""
    readout_kind = type(self.evidence.readout).wire_kind
    rule_kind = type(self.rule).wire_kind
    if readout_kind is None or rule_kind is None:
        return None
    return frozenset({readout_kind}), frozenset({rule_kind})

GateSource

Bases: Protocol

A recipe resolving to a Gate (or None for unconditional) for a model.

Occupies an Intervention's gate slot; Intervention.bind resolves it once. The declared wire kinds are class-level facts so Intervention.wire_kinds() can run before binding; None marks the resolved gating hook-only.

Source code in steerability/algorithms/state_control/common/gating.py

@runtime_checkable
class GateSource(Protocol):
    """A recipe resolving to a `Gate` (or None for unconditional) for a model.

    Occupies an `Intervention`'s gate slot; `Intervention.bind` resolves it once. The declared
    wire kinds are class-level facts so `Intervention.wire_kinds()` can run before binding;
    None marks the resolved gating hook-only.
    """

    wire_readouts: ClassVar[frozenset[str] | None]
    wire_rules: ClassVar[frozenset[str] | None]

    def resolve_gate(self, model, tokenizer, *, layout=None, session=None) -> "Gate | None":
        """Return the resolved gate, or None for unconditional configurations."""
        ...

wire_readouts class-attribute

wire_rules class-attribute

resolve_gate(model, tokenizer, *, layout=None, session=None)

Return the resolved gate, or None for unconditional configurations.

Source code in steerability/algorithms/state_control/common/gating.py

def resolve_gate(self, model, tokenizer, *, layout=None, session=None) -> "Gate | None":
    """Return the resolved gate, or None for unconditional configurations."""
    ...

PerKeyThreshold

Per-layer threshold comparison, combined across layers with any/all.

"ge" opens a layer's vote when its value is at or above the threshold; "le" when at or below. aggregate="any" opens a row when any layer passes for that row; "all" requires every layer.

Parameters:

Name	Type	Description	Default
threshold	float	Score threshold for the per-layer comparison.	required
comparator	Literal['ge', 'le']	"ge" (value >= threshold) or "le" (value <= threshold).	'ge'
aggregate	Literal['any', 'all']	"any" or "all" across layers.	'any'

Source code in steerability/algorithms/state_control/common/gating.py

class PerKeyThreshold:
    """Per-layer threshold comparison, combined across layers with any/all.

    `"ge"` opens a layer's vote when its value is at or above the threshold; `"le"` when at or
    below. `aggregate="any"` opens a row when any layer passes for that row; `"all"` requires
    every layer.

    Args:
        threshold: Score threshold for the per-layer comparison.
        comparator: `"ge"` (value >= threshold) or `"le"` (value <= threshold).
        aggregate: `"any"` or `"all"` across layers.
    """

    wire_kind: ClassVar[str | None] = "per_key_threshold"

    def __init__(
        self,
        threshold: float,
        comparator: Literal["ge", "le"] = "ge",
        aggregate: Literal["any", "all"] = "any",
    ):
        if comparator not in ("ge", "le"):
            raise ValueError(f"comparator must be 'ge' or 'le'; got {comparator!r}.")
        if aggregate not in ("any", "all"):
            raise ValueError(f"aggregate must be 'any' or 'all'; got {aggregate!r}.")
        self.threshold = float(threshold)
        self.comparator = comparator
        self.aggregate = aggregate

    def decide(self, values: Mapping[int, torch.Tensor], num_rows: int) -> torch.BoolTensor:
        """Per-row decision across layers; all-closed when no values have arrived."""
        if not values:
            return torch.zeros(num_rows, dtype=torch.bool)
        stacked = torch.stack(list(values.values()), dim=0)  # [L, num_rows]
        if self.comparator == "ge":
            passed = stacked >= self.threshold
        else:
            passed = stacked <= self.threshold
        if self.aggregate == "any":
            return passed.any(dim=0)
        return passed.all(dim=0)

    def is_complete(self, seen: frozenset[int], expected: frozenset[int]) -> bool:
        """True once every expected layer has reported."""
        return seen >= expected

    def export(self) -> "WireForm | None":
        """The `per_key_threshold` wire form with inline `threshold`, `comparator`, and
        `aggregate` params."""
        from .specs import WireForm

        return WireForm(
            kind="per_key_threshold",
            params={
                "threshold": self.threshold,
                "comparator": self.comparator,
                "aggregate": self.aggregate,
            },
        )

aggregate = aggregate instance-attribute

comparator = comparator instance-attribute

threshold = float(threshold) instance-attribute

wire_kind = 'per_key_threshold' class-attribute

__init__(threshold, comparator='ge', aggregate='any')

Source code in steerability/algorithms/state_control/common/gating.py

def __init__(
    self,
    threshold: float,
    comparator: Literal["ge", "le"] = "ge",
    aggregate: Literal["any", "all"] = "any",
):
    if comparator not in ("ge", "le"):
        raise ValueError(f"comparator must be 'ge' or 'le'; got {comparator!r}.")
    if aggregate not in ("any", "all"):
        raise ValueError(f"aggregate must be 'any' or 'all'; got {aggregate!r}.")
    self.threshold = float(threshold)
    self.comparator = comparator
    self.aggregate = aggregate

decide(values, num_rows)

Per-row decision across layers; all-closed when no values have arrived.

Source code in steerability/algorithms/state_control/common/gating.py

def decide(self, values: Mapping[int, torch.Tensor], num_rows: int) -> torch.BoolTensor:
    """Per-row decision across layers; all-closed when no values have arrived."""
    if not values:
        return torch.zeros(num_rows, dtype=torch.bool)
    stacked = torch.stack(list(values.values()), dim=0)  # [L, num_rows]
    if self.comparator == "ge":
        passed = stacked >= self.threshold
    else:
        passed = stacked <= self.threshold
    if self.aggregate == "any":
        return passed.any(dim=0)
    return passed.all(dim=0)

export()

The per_key_threshold wire form with inline threshold, comparator, and aggregate params.

Source code in steerability/algorithms/state_control/common/gating.py

def export(self) -> "WireForm | None":
    """The `per_key_threshold` wire form with inline `threshold`, `comparator`, and
    `aggregate` params."""
    from .specs import WireForm

    return WireForm(
        kind="per_key_threshold",
        params={
            "threshold": self.threshold,
            "comparator": self.comparator,
            "aggregate": self.aggregate,
        },
    )

is_complete(seen, expected)

True once every expected layer has reported.

Source code in steerability/algorithms/state_control/common/gating.py

def is_complete(self, seen: frozenset[int], expected: frozenset[int]) -> bool:
    """True once every expected layer has reported."""
    return seen >= expected

ProjectedCosineReadout

Projected-cosine score: cosine(pooled, tanh(pooled @ P_layer)) with P = dd^T / (d . d).

[K, H] artifacts use row 0. Rank-one projectors are built lazily and cached per (layer_id, device). Returns zeros for layers absent from the directions. The score is approximately |cos(pooled, d)| (alignment with the line spanned by the direction, erasing which side of the direction a state lies on).

Parameters:

Name	Type	Description	Default
directions	SteeringVector | Mapping[int, Tensor]	Per-layer condition directions, a SteeringVector or Mapping[int, torch.Tensor].	required
location	str | None	The boundary the directions were fitted at, or None to skip the check.	None
model_fingerprint	str | None	The fitted model's identity, or None to skip the check.	None

Reference:

• "Programming Refusal with Conditional Activation Steering" Bruce W. Lee, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Erik Miehling, Pierre Dognin, Manish Nagireddy, Amit Dhurandhar https://arxiv.org/abs/2409.05907

Source code in steerability/algorithms/state_control/common/gating.py

class ProjectedCosineReadout:
    """Projected-cosine score: `cosine(pooled, tanh(pooled @ P_layer))` with
    `P = dd^T / (d . d)`.

    `[K, H]` artifacts use row 0. Rank-one projectors are built lazily and cached per
    `(layer_id, device)`. Returns zeros for layers absent from the directions. The score is
    approximately `|cos(pooled, d)|` (alignment with the line spanned by the direction,
    erasing which side of the direction a state lies on).

    Args:
        directions: Per-layer condition directions, a `SteeringVector` or
            `Mapping[int, torch.Tensor]`.
        location: The boundary the directions were fitted at, or None to skip the check.
        model_fingerprint: The fitted model's identity, or None to skip the check.

    Reference:

    - "Programming Refusal with Conditional Activation Steering"
      Bruce W. Lee, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Erik Miehling, Pierre Dognin,
      Manish Nagireddy, Amit Dhurandhar
      [https://arxiv.org/abs/2409.05907](https://arxiv.org/abs/2409.05907)
    """

    wire_kind: ClassVar[str | None] = "projected_cosine"

    def __init__(
        self,
        directions: SteeringVector | Mapping[int, torch.Tensor],
        location: str | None = None,
        model_fingerprint: str | None = None,
    ):
        self.directions = _extract_directions(directions, type(self).__name__)
        self.location = location
        self.model_fingerprint = model_fingerprint
        self._projector_cache: dict[tuple[int, torch.device], torch.Tensor] = {}

    def _projector(self, layer_id: int, device: torch.device) -> torch.Tensor | None:
        """Cached `[H, H]` rank-one projector for a layer, or None when the layer is absent."""
        key = (layer_id, device)
        cached = self._projector_cache.get(key)
        if cached is None:
            direction = self.directions.get(layer_id)
            if direction is None:
                return None
            cached = rank_one_projector(_feature_row(direction).to(device=device)).to(device)
            self._projector_cache[key] = cached
        return cached

    @torch.no_grad()
    def __call__(self, pooled: torch.Tensor, layer_id: int) -> torch.Tensor:
        """Return `[B]` projected-cosine values for `layer_id` (zeros if the layer is absent)."""
        projector = self._projector(layer_id, pooled.device)
        if projector is None:
            return torch.zeros(pooled.size(0), dtype=torch.float32)
        return projected_cosine_similarity_tensor(pooled, projector.to(pooled.dtype)).float().cpu()

    def export(self, layer_ids: tuple[int, ...]) -> "WireForm | None":
        """The `projected_cosine` wire form: `directions [L, H]` row-aligned with `layer_ids`."""
        from .specs import WireForm

        if any(lid not in self.directions for lid in layer_ids):
            return None
        stacked = torch.stack(
            [_feature_row(self.directions[lid]).reshape(-1).to(torch.float32) for lid in layer_ids]
        )
        return WireForm(kind="projected_cosine", tensors={"directions": stacked})

_projector_cache = {} instance-attribute

directions = _extract_directions(directions, type(self).__name__) instance-attribute

location = location instance-attribute

model_fingerprint = model_fingerprint instance-attribute

wire_kind = 'projected_cosine' class-attribute

__call__(pooled, layer_id)

Return [B] projected-cosine values for layer_id (zeros if the layer is absent).

Source code in steerability/algorithms/state_control/common/gating.py

@torch.no_grad()
def __call__(self, pooled: torch.Tensor, layer_id: int) -> torch.Tensor:
    """Return `[B]` projected-cosine values for `layer_id` (zeros if the layer is absent)."""
    projector = self._projector(layer_id, pooled.device)
    if projector is None:
        return torch.zeros(pooled.size(0), dtype=torch.float32)
    return projected_cosine_similarity_tensor(pooled, projector.to(pooled.dtype)).float().cpu()

__init__(directions, location=None, model_fingerprint=None)

Source code in steerability/algorithms/state_control/common/gating.py

def __init__(
    self,
    directions: SteeringVector | Mapping[int, torch.Tensor],
    location: str | None = None,
    model_fingerprint: str | None = None,
):
    self.directions = _extract_directions(directions, type(self).__name__)
    self.location = location
    self.model_fingerprint = model_fingerprint
    self._projector_cache: dict[tuple[int, torch.device], torch.Tensor] = {}

_projector(layer_id, device)

Cached [H, H] rank-one projector for a layer, or None when the layer is absent.

Source code in steerability/algorithms/state_control/common/gating.py

def _projector(self, layer_id: int, device: torch.device) -> torch.Tensor | None:
    """Cached `[H, H]` rank-one projector for a layer, or None when the layer is absent."""
    key = (layer_id, device)
    cached = self._projector_cache.get(key)
    if cached is None:
        direction = self.directions.get(layer_id)
        if direction is None:
            return None
        cached = rank_one_projector(_feature_row(direction).to(device=device)).to(device)
        self._projector_cache[key] = cached
    return cached

export(layer_ids)

The projected_cosine wire form: directions [L, H] row-aligned with layer_ids.

Source code in steerability/algorithms/state_control/common/gating.py

def export(self, layer_ids: tuple[int, ...]) -> "WireForm | None":
    """The `projected_cosine` wire form: `directions [L, H]` row-aligned with `layer_ids`."""
    from .specs import WireForm

    if any(lid not in self.directions for lid in layer_ids):
        return None
    stacked = torch.stack(
        [_feature_row(self.directions[lid]).reshape(-1).to(torch.float32) for lid in layer_ids]
    )
    return WireForm(kind="projected_cosine", tensors={"directions": stacked})

Readout

Bases: Protocol

Per-layer map from pooled hidden state to a per-row value.

Callable as (pooled [B, H], layer_id) -> Tensor [B] (float32, cpu). Readouts are stateless per-layer functions of pooled hidden states; pooling is declared once, on Evidence. location and model_fingerprint are validated by Intervention.bind when non-None. export(layer_ids) returns the wire form (kind plus row-aligned tensors) or None; wire_kind is the class-level kind name, or None for readouts with no wire form.

Source code in steerability/algorithms/state_control/common/gating.py

@runtime_checkable
class Readout(Protocol):
    """Per-layer map from pooled hidden state to a per-row value.

    Callable as `(pooled [B, H], layer_id) -> Tensor [B]` (float32, cpu). Readouts are
    stateless per-layer functions of pooled hidden states; pooling is declared once, on
    `Evidence`. `location` and `model_fingerprint` are validated by `Intervention.bind` when
    non-None. `export(layer_ids)` returns the wire form (kind plus row-aligned tensors) or
    None; `wire_kind` is the class-level kind name, or None for readouts with no wire form.
    """

    wire_kind: ClassVar[str | None]
    location: str | None
    model_fingerprint: str | None

    def __call__(self, pooled: torch.Tensor, layer_id: int) -> torch.Tensor: ...

    def export(self, layer_ids: tuple[int, ...]) -> "WireForm | None": ...

location instance-attribute

model_fingerprint instance-attribute

wire_kind class-attribute

__call__(pooled, layer_id)

Source code in steerability/algorithms/state_control/common/gating.py

def __call__(self, pooled: torch.Tensor, layer_id: int) -> torch.Tensor: ...

export(layer_ids)

Source code in steerability/algorithms/state_control/common/gating.py

def export(self, layer_ids: tuple[int, ...]) -> "WireForm | None": ...

Rule

Bases: Protocol

Pure decision over per-layer evidence values. Stateless.

decide maps the collected values (one [num_rows] tensor per layer) to a [num_rows] boolean row tensor; is_complete reports whether the collected evidence suffices to decide; export() returns the wire form or None; wire_kind is the class-level kind name, or None for rules with no wire form.

Source code in steerability/algorithms/state_control/common/gating.py

@runtime_checkable
class Rule(Protocol):
    """Pure decision over per-layer evidence values. Stateless.

    `decide` maps the collected values (one `[num_rows]` tensor per layer) to a `[num_rows]`
    boolean row tensor; `is_complete` reports whether the collected evidence suffices to
    decide; `export()` returns the wire form or None; `wire_kind` is the class-level kind
    name, or None for rules with no wire form.
    """

    wire_kind: ClassVar[str | None]

    def decide(self, values: Mapping[int, torch.Tensor], num_rows: int) -> torch.BoolTensor: ...

    def is_complete(self, seen: frozenset[int], expected: frozenset[int]) -> bool: ...

    def export(self) -> "WireForm | None": ...

wire_kind class-attribute

decide(values, num_rows)

Source code in steerability/algorithms/state_control/common/gating.py

def decide(self, values: Mapping[int, torch.Tensor], num_rows: int) -> torch.BoolTensor: ...

export()

Source code in steerability/algorithms/state_control/common/gating.py

def export(self) -> "WireForm | None": ...

is_complete(seen, expected)

Source code in steerability/algorithms/state_control/common/gating.py

def is_complete(self, seen: frozenset[int], expected: frozenset[int]) -> bool: ...

SumThreshold

Open where sum over layers of values + bias >= 0 (ties open).

Parameters:

Name	Type	Description	Default
bias	float	Offset added to the summed values before the zero comparison.	0.0

Source code in steerability/algorithms/state_control/common/gating.py

class SumThreshold:
    """Open where `sum over layers of values + bias >= 0` (ties open).

    Args:
        bias: Offset added to the summed values before the zero comparison.
    """

    wire_kind: ClassVar[str | None] = "sum_threshold"

    def __init__(self, bias: float = 0.0):
        self.bias = float(bias)

    def decide(self, values: Mapping[int, torch.Tensor], num_rows: int) -> torch.BoolTensor:
        """Per-row decision at the bias; all-closed when no values have arrived."""
        if not values:
            return torch.zeros(num_rows, dtype=torch.bool)
        total = torch.stack(list(values.values()), dim=0).sum(dim=0)
        return total + self.bias >= 0

    def is_complete(self, seen: frozenset[int], expected: frozenset[int]) -> bool:
        """True once every expected layer has reported."""
        return seen >= expected

    def export(self) -> "WireForm | None":
        """The `sum_threshold` wire form with the inline `bias` param."""
        from .specs import WireForm

        return WireForm(kind="sum_threshold", params={"bias": self.bias})

bias = float(bias) instance-attribute

wire_kind = 'sum_threshold' class-attribute

__init__(bias=0.0)

Source code in steerability/algorithms/state_control/common/gating.py

def __init__(self, bias: float = 0.0):
    self.bias = float(bias)

decide(values, num_rows)

Per-row decision at the bias; all-closed when no values have arrived.

Source code in steerability/algorithms/state_control/common/gating.py

def decide(self, values: Mapping[int, torch.Tensor], num_rows: int) -> torch.BoolTensor:
    """Per-row decision at the bias; all-closed when no values have arrived."""
    if not values:
        return torch.zeros(num_rows, dtype=torch.bool)
    total = torch.stack(list(values.values()), dim=0).sum(dim=0)
    return total + self.bias >= 0

export()

The sum_threshold wire form with the inline bias param.

Source code in steerability/algorithms/state_control/common/gating.py

def export(self) -> "WireForm | None":
    """The `sum_threshold` wire form with the inline `bias` param."""
    from .specs import WireForm

    return WireForm(kind="sum_threshold", params={"bias": self.bias})

is_complete(seen, expected)

True once every expected layer has reported.

Source code in steerability/algorithms/state_control/common/gating.py

def is_complete(self, seen: frozenset[int], expected: frozenset[int]) -> bool:
    """True once every expected layer has reported."""
    return seen >= expected

_extract_directions(artifact, who)

Validate and extract a concrete per-layer directions mapping from an artifact.

Source code in steerability/algorithms/state_control/common/gating.py

def _extract_directions(
    artifact: SteeringVector | Mapping[int, torch.Tensor], who: str
) -> dict[int, torch.Tensor]:
    """Validate and extract a concrete per-layer directions mapping from an artifact."""
    if isinstance(artifact, SteeringVector):
        directions = dict(artifact.directions)
    elif isinstance(artifact, Mapping):
        directions = dict(artifact)
    else:
        raise TypeError(
            f"{who} expects a concrete SteeringVector or Mapping[int, Tensor]; "
            f"got {type(artifact).__name__}."
        )
    if not all(isinstance(v, torch.Tensor) for v in directions.values()):
        raise TypeError(f"{who} directions must be torch.Tensor values.")
    return {int(lid): direction for lid, direction in directions.items()}

_feature_row(direction)

The [H] feature row of a direction; [K, H] artifacts use row 0.

Source code in steerability/algorithms/state_control/common/gating.py

def _feature_row(direction: torch.Tensor) -> torch.Tensor:
    """The `[H]` feature row of a direction; `[K, H]` artifacts use row 0."""
    if direction.ndim == 2:
        return direction[0]
    return direction

gate_from_probe(probe, *, allow_model_mismatch=False)

A gate reproducing the probe's decision: affine evidence summed at the calibrated bias.

The gate reads the probe's layers at the probe's pooling, scores each layer's pooled state as weights[layer] . pooled, and opens where the sum plus probe.bias is at or above zero (ties open), matching Probe.predict bit-for-bit for the same hidden states (affine pooling commutes with scoring).

Parameters:

Name	Type	Description	Default
probe	Probe	The fitted probe whose decision admits the intervention.	required
allow_model_mismatch	bool	When True, the readout's model_fingerprint is set to None, which disarms the intervention's model-identity check.	False

Returns:

Type	Description
Gate	The assembled Gate.

Source code in steerability/algorithms/state_control/common/gating.py

def gate_from_probe(probe: Probe, *, allow_model_mismatch: bool = False) -> Gate:
    """A gate reproducing the probe's decision: affine evidence summed at the calibrated bias.

    The gate reads the probe's layers at the probe's pooling, scores each layer's pooled state
    as `weights[layer] . pooled`, and opens where the sum plus `probe.bias` is at or above
    zero (ties open), matching `Probe.predict` bit-for-bit for the same hidden states (affine
    pooling commutes with scoring).

    Args:
        probe: The fitted probe whose decision admits the intervention.
        allow_model_mismatch: When True, the readout's `model_fingerprint` is set to None,
            which disarms the intervention's model-identity check.

    Returns:
        The assembled `Gate`.
    """
    readout = AffineReadout(
        dict(probe.weights),
        location=probe.location,
        model_fingerprint=None if allow_model_mismatch else probe.meta.get("model_fingerprint"),
    )
    return Gate(
        Evidence(tuple(probe.layer_ids), readout, pooling=probe.pooling),
        SumThreshold(bias=probe.bias),
    )

projected_cosine_similarity(hidden_state, projector)

Projected-cosine score of one aggregated hidden state.

Projects the hidden state through the condition subspace projector, applies tanh, then computes cosine similarity with the original.

Parameters:

Name	Type	Description	Default
hidden_state	Tensor	Shape [H], an aggregated hidden state.	required
projector	Tensor	Shape [H, H] outer-product projection matrix.	required

Returns:

Type	Description
float	Cosine similarity as a float.

Source code in steerability/algorithms/state_control/common/gating.py

@torch.no_grad()
def projected_cosine_similarity(
    hidden_state: torch.Tensor,
    projector: torch.Tensor,
) -> float:
    """Projected-cosine score of one aggregated hidden state.

    Projects the hidden state through the condition subspace projector, applies tanh, then
    computes cosine similarity with the original.

    Args:
        hidden_state: Shape `[H]`, an aggregated hidden state.
        projector: Shape `[H, H]` outer-product projection matrix.

    Returns:
        Cosine similarity as a float.
    """
    score = projected_cosine_similarity_tensor(hidden_state.unsqueeze(0), projector)[0]
    return float(score.item())

projected_cosine_similarity_tensor(hidden, projector)

Cosine similarity between rows of hidden and their tanh'd projections, one score per row.

Parameters:

Name	Type	Description	Default
hidden	Tensor	Shape [..., H].	required
projector	Tensor	Shape [H, H] outer-product projection matrix.	required

Returns:

Type	Description
Tensor	Scores of shape [...] (float32).

Source code in steerability/algorithms/state_control/common/gating.py

@torch.no_grad()
def projected_cosine_similarity_tensor(
    hidden: torch.Tensor,
    projector: torch.Tensor,
) -> torch.Tensor:
    """Cosine similarity between rows of `hidden` and their tanh'd projections, one score per row.

    Args:
        hidden: Shape `[..., H]`.
        projector: Shape `[H, H]` outer-product projection matrix.

    Returns:
        Scores of shape `[...]` (float32).
    """
    hidden = hidden.float()
    projector = projector.float()
    projected = torch.tanh(hidden @ projector)  # projector is symmetric
    numerator = (hidden * projected).sum(dim=-1)
    denominator = hidden.norm(dim=-1) * projected.norm(dim=-1) + 1e-8
    return numerator / denominator

rank_one_projector(direction)

Build the rank-one projector cc^T / (c^T c) for a direction.

Parameters:

Name	Type	Description	Default
direction	Tensor	Shape [H].	required

Returns:

Type	Description
Tensor	Projection matrix of shape [H, H].

Source code in steerability/algorithms/state_control/common/gating.py

def rank_one_projector(direction: torch.Tensor) -> torch.Tensor:
    """Build the rank-one projector `cc^T / (c^T c)` for a direction.

    Args:
        direction: Shape `[H]`.

    Returns:
        Projection matrix of shape `[H, H]`.
    """
    if direction.ndim != 1:
        raise ValueError(f"direction must be 1-D [H]; got shape {tuple(direction.shape)}.")
    c = direction.float()
    return torch.outer(c, c) / (c @ c + 1e-8)

hook_utils

Utilities for hook registration and model inspection.

extract_hidden_states(input_args, input_kwargs)

Extract hidden_states tensor from a pre-hook's arguments.

HuggingFace transformer layers receive hidden_states either as the first positional argument or as a keyword argument.

Parameters:

Name	Type	Description	Default
input_args	tuple	Positional args from the pre-hook.	required
input_kwargs	dict	Keyword args from the pre-hook.	required

Returns:

Type	Description
Tensor | None	The hidden_states tensor, or None if not found.

Source code in steerability/algorithms/state_control/common/hook_utils.py

def extract_hidden_states(input_args: tuple, input_kwargs: dict) -> torch.Tensor | None:
    """Extract hidden_states tensor from a pre-hook's arguments.

    HuggingFace transformer layers receive hidden_states either as the
    first positional argument or as a keyword argument.

    Args:
        input_args: Positional args from the pre-hook.
        input_kwargs: Keyword args from the pre-hook.

    Returns:
        The hidden_states tensor, or None if not found.
    """
    if input_args:
        return input_args[0]
    return input_kwargs.get("hidden_states")

get_model_layer_list(model)

Return (layer_modules, layer_name_strings) for a HuggingFace model.

Supports the decoder-stack roots model.layers (Llama/Mistral/Qwen/Gemma text), model.language_model.layers (composite multimodal wrappers), and transformer.h (GPT-2), including PEFT-wrapped models.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	A HuggingFace causal LM.	required

Returns:

Type	Description
tuple[list, list[str]]	Tuple of (list of nn.Module layers, list of dotted name strings).

Raises:

Type	Description
ValueError	If model architecture is not recognized.

Source code in steerability/algorithms/state_control/common/hook_utils.py

def get_model_layer_list(model: PreTrainedModel) -> tuple[list, list[str]]:
    """Return (layer_modules, layer_name_strings) for a HuggingFace model.

    Supports the decoder-stack roots `model.layers` (Llama/Mistral/Qwen/Gemma text),
    `model.language_model.layers` (composite multimodal wrappers), and `transformer.h` (GPT-2),
    including PEFT-wrapped models.

    Args:
        model: A HuggingFace causal LM.

    Returns:
        Tuple of (list of nn.Module layers, list of dotted name strings).

    Raises:
        ValueError: If model architecture is not recognized.
    """
    layout = resolve_model_layout(model)
    names = layout.layer_names
    modules = [model.get_submodule(name) for name in names]
    return modules, names

get_norm_module_names(model)

Return (layer_id, module_path) pairs for the per-layer normalization sub-modules.

Returns the residual-stream normalization sub-modules named by the resolved layout:

• gemma_style: input_layernorm, pre_feedforward_layernorm.
• llama_style (Llama/Mistral/Qwen): input_layernorm, post_attention_layernorm.
• gpt2_style: ln_1, ln_2.

Only names that exist on the module are returned, sorted by (layer_id, module_path).

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	A HuggingFace causal LM.	required

Returns:

Type	Description
list[tuple[int, str]]	List of (layer_id, dotted_module_path) pairs, one per normalization sub-module.

Raises:

Type	Description
ValueError	If the model architecture is not recognized.

Source code in steerability/algorithms/state_control/common/hook_utils.py

def get_norm_module_names(model: PreTrainedModel) -> list[tuple[int, str]]:
    """Return (layer_id, module_path) pairs for the per-layer normalization sub-modules.

    Returns the residual-stream normalization sub-modules named by the resolved layout:

    - `gemma_style`: `input_layernorm`, `pre_feedforward_layernorm`.
    - `llama_style` (Llama/Mistral/Qwen): `input_layernorm`, `post_attention_layernorm`.
    - `gpt2_style`: `ln_1`, `ln_2`.

    Only names that exist on the module are returned, sorted by (layer_id, module_path).

    Args:
        model: A HuggingFace causal LM.

    Returns:
        List of `(layer_id, dotted_module_path)` pairs, one per normalization sub-module.

    Raises:
        ValueError: If the model architecture is not recognized.
    """
    layout = resolve_model_layout(model)

    result: list[tuple[int, str]] = []
    for layer_id, layer_name in enumerate(layout.layer_names):
        layer = model.get_submodule(layer_name)
        for attr in layout.norm_attrs:
            if hasattr(layer, attr) and getattr(layer, attr) is not None:
                result.append((layer_id, f"{layer_name}.{attr}"))
    result.sort(key=lambda pair: (pair[0], pair[1]))
    return result

replace_hidden_states(input_args, input_kwargs, new_hidden)

Return modified (input_args, input_kwargs) with hidden_states replaced.

Parameters:

Name	Type	Description	Default
input_args	tuple	Original positional args.	required
input_kwargs	dict	Original keyword args.	required
new_hidden	Tensor	Replacement hidden states tensor.	required

Returns:

Type	Description
tuple[tuple, dict]	Tuple of (new_input_args, new_input_kwargs).

Source code in steerability/algorithms/state_control/common/hook_utils.py

def replace_hidden_states(
    input_args: tuple,
    input_kwargs: dict,
    new_hidden: torch.Tensor,
) -> tuple[tuple, dict]:
    """Return modified (input_args, input_kwargs) with hidden_states replaced.

    Args:
        input_args: Original positional args.
        input_kwargs: Original keyword args.
        new_hidden: Replacement hidden states tensor.

    Returns:
        Tuple of (new_input_args, new_input_kwargs).
    """
    if input_args:
        return (new_hidden, *input_args[1:]), input_kwargs
    input_kwargs = dict(input_kwargs)
    input_kwargs["hidden_states"] = new_hidden
    return input_args, input_kwargs

layout_facts

Structural model facts for steer-time preparation.

State controls consume structural facts (layer count, dtype, hidden size) from the steering session's ModelFacts so preparation works the same whether the steering backend holds a live model or only a layout. Module-path resolution stays out of this module; hook module names are resolved from the module tree at get_hooks() time.

cast_steering_vector(steering_vector, layout)

A clone of steering_vector with per-layer directions cast to the layout dtype.

Device placement is untouched; transforms move tensors to the stream device at apply time.

Parameters:

Name	Type	Description	Default
steering_vector	SteeringVector	The SteeringVector to clone and cast.	required
layout	ModelFacts	The structural layout naming the target dtype.	required

Returns:

Type	Description
SteeringVector	The cast clone.

Source code in steerability/algorithms/state_control/common/layout_facts.py

def cast_steering_vector(steering_vector: SteeringVector, layout: ModelFacts) -> SteeringVector:
    """A clone of `steering_vector` with per-layer directions cast to the layout dtype.

    Device placement is untouched; transforms move tensors to the stream device at apply time.

    Args:
        steering_vector: The `SteeringVector` to clone and cast.
        layout: The structural layout naming the target dtype.

    Returns:
        The cast clone.
    """
    clone = steering_vector.clone()
    dtype = layout_torch_dtype(layout)
    for layer_id, direction in clone.directions.items():
        clone.directions[layer_id] = direction.to(dtype=dtype)
    return clone

layout_torch_dtype(layout)

The torch dtype named by layout.dtype.

Raises:

Type	Description
ValueError	If layout.dtype does not name a torch dtype.

Source code in steerability/algorithms/state_control/common/layout_facts.py

def layout_torch_dtype(layout: ModelFacts) -> torch.dtype:
    """The torch dtype named by `layout.dtype`.

    Raises:
        ValueError: If `layout.dtype` does not name a torch dtype.
    """
    dtype = getattr(torch, layout.dtype, None)
    if not isinstance(dtype, torch.dtype):
        raise ValueError(f"Layout dtype {layout.dtype!r} does not name a torch dtype.")
    return dtype

resolve_layout(model=None, session=None)

Structural facts from the session's layout, else derived from the live model.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	A live model, consulted only when session is None.	None
session	SteeringSession | None	A SteeringSession whose layout property carries the facts.	None

Returns:

Type	Description
ModelFacts	The structural ModelFacts.

Raises:

Type	Description
ValueError	If neither a session nor a model is available.

Source code in steerability/algorithms/state_control/common/layout_facts.py

def resolve_layout(model: PreTrainedModel | None = None, session: SteeringSession | None = None) -> ModelFacts:
    """Structural facts from the session's layout, else derived from the live model.

    Args:
        model: A live model, consulted only when `session` is None.
        session: A `SteeringSession` whose `layout` property carries the facts.

    Returns:
        The structural `ModelFacts`.

    Raises:
        ValueError: If neither a session nor a model is available.
    """
    if session is not None:
        return session.layout
    if model is None:
        raise ValueError(
            "Structural facts require a steering session (session.layout) or a live model; "
            "vector-supplied configurations may steer with model=None only when a session is given."
        )
    _, layer_names = get_model_layer_list(model)
    text_cfg = text_config(model)
    hidden_size = text_cfg.hidden_size
    num_heads = getattr(text_cfg, "num_attention_heads", None)
    head_dim = getattr(text_cfg, "head_dim", None)
    if head_dim is None and num_heads:
        head_dim = hidden_size // num_heads
    return ModelFacts(
        num_layers=len(layer_names),
        hidden_size=hidden_size,
        num_attention_heads=num_heads,
        head_dim=head_dim,
        dtype=str(model.dtype).removeprefix("torch."),
        model_fingerprint=model_fingerprint(model),
        model_type=getattr(model.config, "model_type", None),
        model_ref=getattr(model, "name_or_path", None),
    )

lowering

The wire compiler from bound interventions to intervention-spec payloads.

lower_interventions compiles a bound Intervention tuple into an InterventionSpec for intervention-capable backends, folding each component's wire form and content-addressing tensor payloads via artifact_id_for. Its counterpart, build_hooks in runtime.py, compiles the same IR into torch hooks for one generation.

_lower_gate(gate, boundary, num_layers, register)

The wire gate for one intervention, or None (raised to the caller as hook-only).

A None gate lowers to an ungated op (the caller omits the gate field). The gate's evidence layers are mapped to wire indices at the readout's declared boundary (falling back to the intervention's boundary when the readout declares none), preserving the readout tensor's row alignment since the mapping is order-preserving.

Source code in steerability/algorithms/state_control/common/lowering.py

def _lower_gate(gate: Gate | None, boundary: Boundary, num_layers: int, register) -> dict[str, Any] | None:
    """The wire gate for one intervention, or None (raised to the caller as hook-only).

    A None gate lowers to an ungated op (the caller omits the gate field). The gate's evidence
    layers are mapped to wire indices at the readout's declared boundary (falling back to the
    intervention's boundary when the readout declares none), preserving the readout tensor's
    row alignment since the mapping is order-preserving.
    """
    wire = gate.export(register)
    if wire is None:
        return None
    condition_boundary = gate.evidence.readout.location or boundary
    mapped = _map_condition_layers(wire["layers"], condition_boundary, num_layers)
    if mapped is None:
        return None
    wire["layers"] = mapped
    return wire

_map_behavior_layer(layer_id, boundary, site, num_layers)

Map a toolkit behavior layer onto its wire layer index, or None when unmappable.

A wire op applies at the residual-stream boundary after decoder layer N. A "layer_output" hook at layer l is wire layer l; a "layer_input" hook at layer l is wire layer l - 1 (layer 0 has no wire form); the "o_proj" site keeps its layer index, matching the wire head_additive placement.

Source code in steerability/algorithms/state_control/common/lowering.py

def _map_behavior_layer(layer_id: int, boundary: Boundary, site: Site, num_layers: int) -> int | None:
    """Map a toolkit behavior layer onto its wire layer index, or None when unmappable.

    A wire op applies at the residual-stream boundary after decoder layer `N`. A
    `"layer_output"` hook at layer `l` is wire layer `l`; a `"layer_input"` hook at layer `l`
    is wire layer `l - 1` (layer 0 has no wire form); the `"o_proj"` site keeps its layer
    index, matching the wire `head_additive` placement.
    """
    if site == "o_proj":
        mapped = layer_id
    elif boundary == "layer_input":
        mapped = layer_id - 1
    else:
        mapped = layer_id
    return mapped if 0 <= mapped < num_layers else None

_map_condition_layers(layer_ids, boundary, num_layers)

Map toolkit condition layers onto wire layer indices, or None when unmappable.

A wire gate's condition layers read the materialized input of decoder layer N, so layers read at "layer_output" shift to l + 1.

Source code in steerability/algorithms/state_control/common/lowering.py

def _map_condition_layers(
    layer_ids: Sequence[int], boundary: Boundary, num_layers: int
) -> list[int] | None:
    """Map toolkit condition layers onto wire layer indices, or None when unmappable.

    A wire gate's condition layers read the materialized input of decoder layer `N`, so
    layers read at `"layer_output"` shift to `l + 1`.
    """
    offset = 1 if boundary == "layer_output" else 0
    mapped = [int(layer_id) + offset for layer_id in layer_ids]
    if all(0 <= layer_id < num_layers for layer_id in mapped):
        return mapped
    return None

artifact_id_for(tensors)

The content-addressed artifact id and prepared tensors for a tensor payload.

Tensors are prepared as float32, contiguous, CPU copies (cloned before the cast, so the live steering artifacts are never mutated or aliased), and the id is the SHA-256 over the safetensors serialization with sorted tensor names, matching the plugin registry's write byte-for-byte. Identical logical content therefore yields identical ids regardless of the producing device or dtype.

Parameters:

Name	Type	Description	Default
tensors	Mapping[str, Tensor]	Mapping from tensor name to tensor.	required

Returns:

Type	Description
tuple[str, dict[str, Tensor]]	The sha256:<hex> id and the prepared name-to-tensor mapping.

Source code in steerability/algorithms/state_control/common/lowering.py

def artifact_id_for(tensors: Mapping[str, torch.Tensor]) -> tuple[str, dict[str, torch.Tensor]]:
    """The content-addressed artifact id and prepared tensors for a tensor payload.

    Tensors are prepared as float32, contiguous, CPU copies (cloned before the cast, so the
    live steering artifacts are never mutated or aliased), and the id is the SHA-256 over the
    safetensors serialization with sorted tensor names, matching the plugin registry's `write`
    byte-for-byte. Identical logical content therefore yields identical ids regardless of the
    producing device or dtype.

    Args:
        tensors: Mapping from tensor name to tensor.

    Returns:
        The `sha256:<hex>` id and the prepared name-to-tensor mapping.
    """
    import safetensors.torch

    prepared = {
        name: tensor.detach().to(device="cpu", dtype=torch.float32, copy=True).contiguous()
        for name, tensor in tensors.items()
    }
    data = safetensors.torch.save({name: prepared[name] for name in sorted(prepared)})
    return "sha256:" + hashlib.sha256(data).hexdigest(), prepared

lower_interventions(interventions, *, num_layers)

Lower bound interventions to an InterventionSpec, or None when any element has no wire form.

Folds each component's export, unwrap_modifiers, the scope export, and the gate's export. One wire op is emitted per (intervention, layer), in intervention order then ascending layer order; artifact ids are content hashes, so layers sharing a tensor share one artifact. The assembled spec is pre-flight validated with the plugin's parse_intervention_spec, so a malformed spec fails here with the same E_* code and JSON path the server would return.

Parameters:

Name	Type	Description	Default
interventions	Sequence[Intervention]	Bound interventions, in application order.	required
num_layers	int	Decoder layer count from the model layout.	required

Returns:

Type	Description
'InterventionSpec | None'	The validated spec with tensor payloads attached, or None.

Raises:

Type	Description
ValueError	If the assembled spec fails pre-flight validation (a toolkit-side serialization bug; the message carries the E_* code and JSON path).
ModuleNotFoundError	If vllm_hook_plugins is not installed.

Source code in steerability/algorithms/state_control/common/lowering.py

def lower_interventions(
    interventions: Sequence[Intervention],
    *,
    num_layers: int,
) -> "InterventionSpec | None":
    """Lower bound interventions to an `InterventionSpec`, or None when any element has no
    wire form.

    Folds each component's `export`, `unwrap_modifiers`, the scope export, and the gate's
    `export`. One wire op is emitted per (intervention, layer), in intervention order then
    ascending layer order; artifact ids are content hashes, so layers sharing a tensor share
    one artifact. The assembled spec is pre-flight validated with the plugin's
    `parse_intervention_spec`, so a malformed spec fails here with the same `E_*` code and
    JSON path the server would return.

    Args:
        interventions: Bound interventions, in application order.
        num_layers: Decoder layer count from the model layout.

    Returns:
        The validated spec with tensor payloads attached, or None.

    Raises:
        ValueError: If the assembled spec fails pre-flight validation (a toolkit-side
            serialization bug; the message carries the `E_*` code and JSON path).
        ModuleNotFoundError: If `vllm_hook_plugins` is not installed.
    """
    from steerability.algorithms.core.execution.payloads import InterventionSpec
    from steerability.utils.optional import require

    from .transforms.base import unwrap_modifiers

    schema = require("vllm_hook_plugins.core.schema")

    artifacts: dict[str, dict[str, torch.Tensor]] = {}

    def register(tensors: Mapping[str, torch.Tensor]) -> str:
        artifact_id, prepared = artifact_id_for(tensors)
        artifacts.setdefault(artifact_id, prepared)
        return artifact_id

    ops: list[dict[str, Any]] = []
    for intervention in interventions:
        if not isinstance(intervention.layers, tuple):
            raise ValueError("lower_interventions requires bound interventions; call bind() first.")
        site = intervention.resolved_site()
        if site == "norm_input":
            return None
        scope_wire = intervention.scope.export()
        scope: dict[str, Any] = {"kind": scope_wire.kind, **scope_wire.params}

        gate = intervention.gate
        if gate is not None and not isinstance(gate, Gate):
            raise ValueError("lower_interventions requires a resolved gate; call bind() first.")
        gate_wire: dict[str, Any] | None = None
        if gate is not None:
            gate_wire = _lower_gate(gate, intervention.boundary, num_layers, register)
            if gate_wire is None:
                return None

        core, wrappers = unwrap_modifiers(intervention.transform)
        for layer_id in sorted(intervention.layers):
            form = core.export(layer_id)
            if form is None:
                return None
            wire_layer = _map_behavior_layer(layer_id, intervention.boundary, site, num_layers)
            if wire_layer is None:
                return None
            transform_wire: dict[str, Any] = {"kind": form.kind, **form.params}
            modifier_wires: list[dict[str, Any]] = []
            for wrapper in wrappers:
                if wrapper.modifier_wire_kind(form.kind) is None:
                    return None
                modifier_form = wrapper.export_modifier(layer_id)
                if modifier_form is None:
                    continue  # this wrapper contributes no modifier at this layer
                modifier_wire: dict[str, Any] = {"kind": modifier_form.kind, **modifier_form.params}
                if modifier_form.tensors:
                    modifier_wire["artifact"] = register(modifier_form.tensors)
                modifier_wires.append(modifier_wire)
            transform_wire["modifiers"] = modifier_wires
            if form.tensors:
                transform_wire["artifact"] = register(form.tensors)
            ops.append({
                "layers": [wire_layer],
                "transform": transform_wire,
                "scope": dict(scope),
                "gate": dict(gate_wire) if gate_wire is not None else None,
            })

    if not ops:
        return None

    spec = InterventionSpec(ops=tuple(ops), artifacts=artifacts)
    schema.parse_intervention_spec(spec.to_wire(), num_layers=num_layers)
    return spec

model_layout

Compatibility re-export; the layout lives in steerability.algorithms.core.internals.model_layout.

runtime

Shared hook runtime for transform-based state controls.

TransformHookRuntime builds the hook closures used by residual-stream state controls and owns one control's per-generation mutable state. It covers hidden-state extraction and re-wrapping, KV-cache position tracking, token-scope masking, condition scoring, and gated transform application. Four behaviors define the runtime's contract:

1. Position tracking: Each pass's absolute position offset is read from the hooked module's pass positions when it supplies them (from the cache_position kwarg, or reconstructed from position_ids; decoder layers receive position_ids on every pass), so positions are exact per forwarded sequence even when a decoding driver issues several generate calls or an output control forwards the model mid-step. Hook points whose modules receive neither kwarg (attention output projections, norm sub-modules) fall back to counting, which assumes the model processes the full prompt on the prefill pass and one new token per decode pass; exactly one designated pass-opener hook advances the shared offset by the observed sequence length once per pass, and every other hook in that pass reads the opener's snapshot. The fallback assumes a single generate call per generation.

2. Row gating: Gates hold one decision per logical row, one per prompt. HuggingFace generate may expand the hidden batch to B_logical * num_beams via repeat_interleave. The runtime collapses condition scores from the expanded batch to logical rows before calling gate.update(), and expands gate.open_rows() back to the hidden batch before masking hidden states. Beam siblings of one prompt share that prompt's decision, and each prompt in a batch is gated independently.

3. Condition scoring: A condition hook computes evidence values only while its gate has not frozen its decision. Scoring stops for the remainder of the generation once gate.is_ready() returns True.

4. Auxiliary passes: Forwards marked via auxiliary_pass() (same-model candidate scoring, variant-prompt branches) never feed condition scorers or gates and never advance the fallback counter. Trajectory-aligned auxiliary passes are transformed at their true positions when the pass positions are available; detached ones are never transformed.

HookPoint = Literal['layer_output', 'layer_input'] module-attribute

TransformHookRuntime

Builds hook closures and owns the per-generation position/prefill/mask state for one control.

Parameters:

Name	Type	Description	Default
hook_point	HookPoint	Where the control intervenes. "layer_output" builds forward hooks on the layer output (residual stream after the layer). "layer_input" builds forward pre-hooks; hidden states are extracted from the module's inputs via extract_hidden_states and re-injected via replace_hidden_states. Valid for any module receiving hidden_states as its first positional argument or as the hidden_states= kwarg, such as decoder layers, attention output projections (o_proj/c_proj), and per-layer norm sub-modules.	'layer_output'

Source code in steerability/algorithms/state_control/common/runtime.py

class TransformHookRuntime:
    """Builds hook closures and owns the per-generation position/prefill/mask state for one control.

    Args:
        hook_point: Where the control intervenes. ``"layer_output"`` builds forward hooks on the
            layer output (residual stream after the layer). ``"layer_input"`` builds forward
            pre-hooks; hidden states are extracted from the module's inputs via
            `extract_hidden_states` and re-injected via `replace_hidden_states`. Valid for any
            module receiving `hidden_states` as its first positional argument or as the
            ``hidden_states=`` kwarg, such as decoder layers, attention output projections
            (`o_proj`/`c_proj`), and per-layer norm sub-modules.
    """

    def __init__(self, *, hook_point: HookPoint = "layer_output"):
        if hook_point not in ("layer_output", "layer_input"):
            raise ValueError(f"hook_point must be 'layer_output' or 'layer_input'; got {hook_point!r}.")
        self.hook_point = hook_point

        # per-generation state (set/cleared by reset)
        self._prompt_lens: torch.LongTensor | None = None
        self._prompt_mask: torch.BoolTensor | None = None
        self._offset: int = 0
        self._pass_offset: int = 0
        self._prefill_seen: bool = False
        self._opener_built: bool = False
        self._clock_seen: bool = False
        self._warned: set[str] = set()

    def reset(
        self,
        prompt_lens: torch.LongTensor,
        prompt_mask: torch.Tensor | None = None,
    ) -> None:
        """Clear position/prefill state and store the prompt lengths/mask for this generation.

        Args:
            prompt_lens: Per-row prompt lengths of shape ``[B_logical]`` (from
                `compute_prompt_lens`). Defines the logical batch size for row gating.
            prompt_mask: Optional pad-aware prompt attention mask of shape
                ``[B_logical, T_prompt]`` (True/1 at real tokens). Forwarded to condition scorers
                on the prefill pass so condition scores align with the real (non-pad) prompt
                positions, matching the mask the selector calibrated on.
        """
        self._prompt_lens = prompt_lens
        if prompt_mask is not None:
            mask = torch.as_tensor(prompt_mask)
            if mask.ndim == 1:
                mask = mask.unsqueeze(0)
            if mask.size(0) != prompt_lens.size(0):
                raise ValueError(
                    f"prompt_mask has {mask.size(0)} rows but prompt_lens has "
                    f"{prompt_lens.size(0)}; these must describe the same logical batch."
                )
            self._prompt_mask = mask.bool()
        else:
            self._prompt_mask = None
        self._offset = 0
        self._pass_offset = 0
        self._prefill_seen = False
        self._opener_built = False
        self._clock_seen = False
        self._warned = set()

    @property
    def num_logical_rows(self) -> int:
        """Logical batch size (one row per prompt); 0 before `reset`."""
        return 0 if self._prompt_lens is None else int(self._prompt_lens.size(0))

    def _claim_opener(self, is_pass_opener: bool) -> None:
        """Enforce that at most one hook per generation is designated the pass opener.

        Two openers would advance the shared offset twice per forward pass, silently skewing
        every position-dependent token scope (e.g. `after_prompt` would steer the whole
        prompt). Controls must designate exactly one opener; when a layer hosts both a
        condition and a behavior hook, the one registered first opens the pass. The opener only
        matters for the fallback counter; when positions come from `cache_position`, every hook
        resolves its offset independently.
        """
        if not is_pass_opener:
            return
        if self._opener_built:
            raise ValueError(
                "A pass-opener hook was already built for this generation; exactly one hook "
                "may advance the position offset. When a layer hosts both a condition and a "
                "behavior hook, designate only the first-registered one as the opener."
            )
        self._opener_built = True

    @staticmethod
    def _extract_pass_positions(forward_kwargs: dict | None, seq_len: int) -> torch.Tensor | None:
        """The pass's cache-axis positions from the hooked module's kwargs, when recoverable.

        Prefers the `cache_position` kwarg (transformers threads it into decoder layers only when
        the model caller passes it explicitly). When absent, the padded-cache-axis offset is
        reconstructed from `position_ids`, which transformers v5 forwards to decoder layers on
        every pass: the longest row of a batch is unpadded, so its positions coincide with the
        padded cache axis and `max(position_ids) - (seq_len - 1)` is the pass's start offset.
        Neither kwarg reaches non-decoder hook points (e.g. `o_proj` inputs), which keep the
        counting fallback.
        """
        if not forward_kwargs:
            return None
        positions = forward_kwargs.get("cache_position")
        if positions is not None and torch.is_tensor(positions) and positions.numel() > 0:
            return positions
        position_ids = forward_kwargs.get("position_ids")
        if position_ids is not None and torch.is_tensor(position_ids) and position_ids.numel() > 0:
            offset = int(position_ids.max().item()) - (seq_len - 1)
            if offset >= 0:
                return torch.tensor([offset], device=position_ids.device)
        return None

    def _warn_once(self, key: str, message: str) -> None:
        """Emit `message` as a UserWarning at most once per generation (keyed by `key`)."""
        if key in self._warned:
            return
        self._warned.add(key)
        warnings.warn(message, UserWarning)

    def _position_offset(
        self, seq_len: int, cache_position: torch.Tensor | None, is_pass_opener: bool
    ) -> int | None:
        """Resolve the absolute position offset for the current pass, or None to skip the pass.

        Auxiliary passes (marked via `auxiliary_pass()`) never advance the fallback counter. An
        aligned auxiliary pass is positioned by the pass positions when the hooked module supplies
        them (from `cache_position` or `position_ids`) and skipped otherwise; a detached auxiliary
        pass is always skipped. Ordinary passes take their offset from the pass positions when
        present. The opener maintains the fallback counter on every ordinary pass, so hook points
        without either kwarg keep the one-pass-one-step accounting unchanged and an anomalous pass
        missing them degrades to counting.

        Args:
            seq_len: The sequence length seen by this hook on this call.
            cache_position: The pass positions (from `cache_position` or reconstructed from
                `position_ids`), when the hooked module supplies them.
            is_pass_opener: Whether this hook is the designated pass opener.

        Returns:
            The absolute position offset to use when building the token mask, or None when the
            pass must be skipped.

        Warns:
            UserWarning: Once per generation for each of: an aligned auxiliary pass at a hook point
                without pass positions (its transform is skipped); a multi-token ordinary pass
                after prefill at such a hook point (a multi-call decode pattern that counting cannot
                place); pass positions disappearing after having been observed.
        """
        aux = current_auxiliary_pass()
        if aux is not None:
            if not aux.aligned:
                return None
            if cache_position is not None:
                return int(cache_position[0])
            self._warn_once(
                "aux_without_cache_position",
                "Auxiliary same-model passes cannot be position-mapped at this hook point (the "
                "hooked module receives neither `cache_position` nor `position_ids`); their "
                "transforms are skipped. Hook decoder layers for exact composition with same-model "
                "output controls.",
            )
            return None

        if is_pass_opener:
            if self._prefill_seen:
                if seq_len > 1 and cache_position is None:
                    self._warn_once(
                        "multi_call_without_cache_position",
                        "Multiple generate calls detected at a hook point that receives neither "
                        "`cache_position` nor `position_ids`; position scoping may be skewed for "
                        "this generation.",
                    )
                self._pass_offset = self._offset
                self._offset += seq_len
            else:
                self._pass_offset = 0
                self._offset = seq_len
                self._prefill_seen = True

        if cache_position is not None:
            self._clock_seen = True
            return int(cache_position[0])
        if self._clock_seen:
            self._warn_once(
                "inconsistent_cache_position",
                "Pass positions were available on earlier passes but missing on this one; "
                "falling back to pass counting for this pass.",
            )
        return self._pass_offset

    def _collapse_to_rows(self, values: torch.Tensor | float, hidden_batch: int) -> torch.Tensor | float:
        """Collapse per-hidden-row evidence values down to the logical rows the gate holds.

        Beam search expands the batch via `repeat_interleave` (`[i0, i0, i1, i1]`), so the first
        member of each group represents its logical row, and that row is taken as the group's
        value. Values already at logical size pass through; a bare float passes through for the
        gate to validate (accepted only when `num_rows == 1`).
        """
        if isinstance(values, (int, float)):
            return values
        rows = self.num_logical_rows
        t = torch.as_tensor(values).squeeze()
        if t.ndim > 1:
            raise ValueError(
                f"Gate readout returned a tensor of shape {tuple(torch.as_tensor(values).shape)}; "
                f"expected per-row values of shape [B] (extra dimensions must be size 1)."
            )
        flat = t.reshape(-1)
        if flat.numel() == rows:
            return flat
        if flat.numel() == hidden_batch and rows and hidden_batch % rows == 0:
            factor = hidden_batch // rows
            return flat[::factor]
        raise ValueError(
            f"Gate readout returned {flat.numel()} values for a hidden batch of "
            f"{hidden_batch} and {rows} logical row(s); return one value per hidden row or per "
            f"logical row."
        )

    def _row_mask_for(self, gate: Gate, hidden: torch.Tensor) -> torch.BoolTensor | None:
        """Per-row gate decision expanded to the hidden batch as a `[B_hidden, 1]` mask.

        Returns None when every row is closed (caller short-circuits). `align_mask_to_batch`
        performs the beam expansion and validates divisibility.
        """
        open_rows = gate.open_rows()
        if not bool(open_rows.any()):
            return None
        row_mask = align_mask_to_batch(open_rows.unsqueeze(1), hidden.size(0))  # [B_hidden, 1]
        return row_mask.to(hidden.device)

    def _prefill_prompt_mask(self, hidden: torch.Tensor, pass_offset: int) -> torch.Tensor | None:
        """The stored prompt mask aligned to the hidden batch, on the prefill pass only.

        The first pass may be longer than the prompt when a teacher-forced continuation is
        appended (e.g. `compute_logprobs` forwards `[prompt; ref]` in one pass). The continuation
        columns are not prompt, so the mask is extended with False there, and condition
        aggregation then covers exactly the real prompt tokens, reproducing generation-time
        scoring. A first pass shorter than the stored mask indicates misuse and raises.
        """
        if self._prompt_mask is None or pass_offset != 0:
            return None
        width = self._prompt_mask.size(1)
        seq_len = hidden.size(1)
        if width > seq_len:
            raise ValueError(
                f"Prompt mask length {width} exceeds prefill sequence length {seq_len}."
            )
        mask = self._prompt_mask
        if width < seq_len:  # teacher-forced continuation appended after the prompt
            pad = torch.zeros(mask.size(0), seq_len - width, dtype=torch.bool)
            mask = torch.cat([mask, pad], dim=1)
        mask = align_mask_to_batch(mask, hidden.size(0))
        return mask.to(hidden.device)

    def build_behavior_hook(
        self,
        *,
        layer_id: int,
        transform: BaseTransform,
        gate: Gate | None,
        token_scope: ScopeKind,
        last_k: int | None = None,
        from_position: int | None = None,
        is_pass_opener: bool = False,
        hook_point: HookPoint | None = None,
    ) -> Callable:
        """Build a hook that applies `transform` to the residual stream at `layer_id`, gated by `gate`.

        The transform fires at the intersection of the token-scope mask and the gate's per-row
        decision (expanded across beams); a fully closed gate is a no-op, and a None gate
        leaves every row open.

        Args:
            layer_id: Index of the hooked layer (used to index per-layer transform artifacts).
            transform: The transform to apply at masked positions of open rows.
            gate: Gate consulted per call, or None for unconditional application; row `r` of
                the hidden batch fires only when the gate's logical row `r // beam_factor` is
                open.
            token_scope: Which positions to steer (see `make_token_mask`).
            last_k: Required when `token_scope == "last_k"`.
            from_position: Required when `token_scope == "from_position"`.
            is_pass_opener: Whether this hook advances the shared position offset.
            hook_point: Per-hook boundary override; defaults to the runtime's constructor
                value, so one runtime can host hooks at both boundaries.

        Returns:
            A hook callable suitable for the effective hook point (a forward hook for
            ``"layer_output"``, a forward pre-hook for ``"layer_input"``).
        """
        self._claim_opener(is_pass_opener)
        if (hook_point or self.hook_point) == "layer_output":

            def _forward_hook(module, args, kwargs, output):
                hidden = output[0] if isinstance(output, tuple) else output
                if hidden is None:
                    return output
                hidden = self._apply(hidden, layer_id, transform, gate, token_scope, last_k,
                                     from_position, is_pass_opener, forward_kwargs=kwargs)
                return (hidden,) + output[1:] if isinstance(output, tuple) else hidden

            return _forward_hook

        def _pre_hook(module, input_args, input_kwargs):
            hidden = extract_hidden_states(input_args, input_kwargs)
            if hidden is None:
                return input_args, input_kwargs
            hidden = self._apply(hidden, layer_id, transform, gate, token_scope, last_k,
                                 from_position, is_pass_opener, forward_kwargs=input_kwargs)
            return replace_hidden_states(input_args, input_kwargs, hidden)

        return _pre_hook

    def build_condition_hook(
        self,
        *,
        layer_id: int,
        gate: Gate,
        is_pass_opener: bool = False,
        hook_point: HookPoint | None = None,
    ) -> Callable:
        """Build a read-only hook that scores the residual stream at `layer_id` and updates `gate`.

        The hook never modifies hidden states. On each pass where the gate has not frozen its
        decision (`not gate.is_ready()`), it pools the hidden states per the gate's evidence
        (passing the pad-aware prompt mask on the prefill pass), computes per-row values via
        the evidence readout, collapses beam-expanded values to logical rows, and calls
        `gate.update(rows, key=layer_id)`. Once the gate is ready, scoring is skipped
        entirely, and a gate whose rule never reports complete keeps re-scoring every pass.
        The hook still participates in pass-opener bookkeeping when the lowest hooked layer is
        a condition layer. Auxiliary passes are ignored entirely: no scoring, no gate update,
        no accounting.

        Args:
            layer_id: Index of the hooked layer.
            gate: Gate whose evidence is read at this layer and fed the per-row values.
            is_pass_opener: Whether this hook advances the shared position offset.
            hook_point: Per-hook boundary override; defaults to the runtime's constructor
                value.

        Returns:
            A hook callable suitable for the effective hook point.
        """
        self._claim_opener(is_pass_opener)

        def _score(hidden: torch.Tensor, forward_kwargs: dict | None) -> None:
            if current_auxiliary_pass() is not None:
                return
            cache_position = self._extract_pass_positions(forward_kwargs, hidden.size(1))
            pass_offset = self._position_offset(hidden.size(1), cache_position, is_pass_opener)
            if gate.is_ready():
                return
            prompt_mask = self._prefill_prompt_mask(hidden, pass_offset)
            pooled = aggregate_condition_hidden(
                hidden, gate.evidence.pooling, attention_mask=prompt_mask
            )
            values = gate.evidence.readout(pooled, layer_id)
            gate.update(self._collapse_to_rows(values, hidden.size(0)), key=layer_id)

        if (hook_point or self.hook_point) == "layer_output":

            def _forward_hook(module, args, kwargs, output):
                hidden = output[0] if isinstance(output, tuple) else output
                if hidden is None:
                    return output
                _score(hidden, kwargs)
                return output

            return _forward_hook

        def _pre_hook(module, input_args, input_kwargs):
            hidden = extract_hidden_states(input_args, input_kwargs)
            if hidden is None:
                return input_args, input_kwargs
            _score(hidden, input_kwargs)
            return input_args, input_kwargs

        return _pre_hook

    def _apply(
        self,
        hidden: torch.Tensor,
        layer_id: int,
        transform: BaseTransform,
        gate: Gate | None,
        token_scope: ScopeKind,
        last_k: int | None,
        from_position: int | None,
        is_pass_opener: bool,
        forward_kwargs: dict | None = None,
    ) -> torch.Tensor:
        """Mask the current pass by token scope and per-row gate decision, then apply the transform.

        The pass's absolute position offset is forwarded to `transform.apply` as
        `position_offset`, so position-dependent transforms place their edits in absolute
        sequence coordinates. A None gate leaves every row open. Auxiliary passes without a
        resolvable position are returned unchanged.
        """
        seq_len = hidden.size(1)
        cache_position = self._extract_pass_positions(forward_kwargs, seq_len)
        pass_offset = self._position_offset(seq_len, cache_position, is_pass_opener)
        if pass_offset is None:
            return hidden

        row_mask = None
        if gate is not None:
            row_mask = self._row_mask_for(gate, hidden)  # [B_hidden, 1] or None (all closed)
            if row_mask is None:
                return hidden

        mask = make_token_mask(
            token_scope,
            seq_len=seq_len,
            prompt_lens=self._prompt_lens.to(hidden.device),
            last_k=last_k,
            from_position=from_position,
            position_offset=pass_offset,
        )
        mask = align_mask_to_batch(mask, hidden.size(0))  # beam search expands the batch
        if row_mask is not None:
            mask = mask & row_mask
        if not bool(mask.any()):
            return hidden
        return transform.apply(hidden, layer_id=layer_id, token_mask=mask, position_offset=pass_offset)

_clock_seen = False instance-attribute

_offset = 0 instance-attribute

_opener_built = False instance-attribute

_pass_offset = 0 instance-attribute

_prefill_seen = False instance-attribute

_prompt_lens = None instance-attribute

_prompt_mask = None instance-attribute

_warned = set() instance-attribute

hook_point = hook_point instance-attribute

num_logical_rows property

Logical batch size (one row per prompt); 0 before reset.

__init__(*, hook_point='layer_output')

Source code in steerability/algorithms/state_control/common/runtime.py

def __init__(self, *, hook_point: HookPoint = "layer_output"):
    if hook_point not in ("layer_output", "layer_input"):
        raise ValueError(f"hook_point must be 'layer_output' or 'layer_input'; got {hook_point!r}.")
    self.hook_point = hook_point

    # per-generation state (set/cleared by reset)
    self._prompt_lens: torch.LongTensor | None = None
    self._prompt_mask: torch.BoolTensor | None = None
    self._offset: int = 0
    self._pass_offset: int = 0
    self._prefill_seen: bool = False
    self._opener_built: bool = False
    self._clock_seen: bool = False
    self._warned: set[str] = set()

_apply(hidden, layer_id, transform, gate, token_scope, last_k, from_position, is_pass_opener, forward_kwargs=None)

Mask the current pass by token scope and per-row gate decision, then apply the transform.

The pass's absolute position offset is forwarded to transform.apply as position_offset, so position-dependent transforms place their edits in absolute sequence coordinates. A None gate leaves every row open. Auxiliary passes without a resolvable position are returned unchanged.

Source code in steerability/algorithms/state_control/common/runtime.py

def _apply(
    self,
    hidden: torch.Tensor,
    layer_id: int,
    transform: BaseTransform,
    gate: Gate | None,
    token_scope: ScopeKind,
    last_k: int | None,
    from_position: int | None,
    is_pass_opener: bool,
    forward_kwargs: dict | None = None,
) -> torch.Tensor:
    """Mask the current pass by token scope and per-row gate decision, then apply the transform.

    The pass's absolute position offset is forwarded to `transform.apply` as
    `position_offset`, so position-dependent transforms place their edits in absolute
    sequence coordinates. A None gate leaves every row open. Auxiliary passes without a
    resolvable position are returned unchanged.
    """
    seq_len = hidden.size(1)
    cache_position = self._extract_pass_positions(forward_kwargs, seq_len)
    pass_offset = self._position_offset(seq_len, cache_position, is_pass_opener)
    if pass_offset is None:
        return hidden

    row_mask = None
    if gate is not None:
        row_mask = self._row_mask_for(gate, hidden)  # [B_hidden, 1] or None (all closed)
        if row_mask is None:
            return hidden

    mask = make_token_mask(
        token_scope,
        seq_len=seq_len,
        prompt_lens=self._prompt_lens.to(hidden.device),
        last_k=last_k,
        from_position=from_position,
        position_offset=pass_offset,
    )
    mask = align_mask_to_batch(mask, hidden.size(0))  # beam search expands the batch
    if row_mask is not None:
        mask = mask & row_mask
    if not bool(mask.any()):
        return hidden
    return transform.apply(hidden, layer_id=layer_id, token_mask=mask, position_offset=pass_offset)

_claim_opener(is_pass_opener)

Enforce that at most one hook per generation is designated the pass opener.

Two openers would advance the shared offset twice per forward pass, silently skewing every position-dependent token scope (e.g. after_prompt would steer the whole prompt). Controls must designate exactly one opener; when a layer hosts both a condition and a behavior hook, the one registered first opens the pass. The opener only matters for the fallback counter; when positions come from cache_position, every hook resolves its offset independently.

Source code in steerability/algorithms/state_control/common/runtime.py

def _claim_opener(self, is_pass_opener: bool) -> None:
    """Enforce that at most one hook per generation is designated the pass opener.

    Two openers would advance the shared offset twice per forward pass, silently skewing
    every position-dependent token scope (e.g. `after_prompt` would steer the whole
    prompt). Controls must designate exactly one opener; when a layer hosts both a
    condition and a behavior hook, the one registered first opens the pass. The opener only
    matters for the fallback counter; when positions come from `cache_position`, every hook
    resolves its offset independently.
    """
    if not is_pass_opener:
        return
    if self._opener_built:
        raise ValueError(
            "A pass-opener hook was already built for this generation; exactly one hook "
            "may advance the position offset. When a layer hosts both a condition and a "
            "behavior hook, designate only the first-registered one as the opener."
        )
    self._opener_built = True

_collapse_to_rows(values, hidden_batch)

Collapse per-hidden-row evidence values down to the logical rows the gate holds.

Beam search expands the batch via repeat_interleave ([i0, i0, i1, i1]), so the first member of each group represents its logical row, and that row is taken as the group's value. Values already at logical size pass through; a bare float passes through for the gate to validate (accepted only when num_rows == 1).

Source code in steerability/algorithms/state_control/common/runtime.py

def _collapse_to_rows(self, values: torch.Tensor | float, hidden_batch: int) -> torch.Tensor | float:
    """Collapse per-hidden-row evidence values down to the logical rows the gate holds.

    Beam search expands the batch via `repeat_interleave` (`[i0, i0, i1, i1]`), so the first
    member of each group represents its logical row, and that row is taken as the group's
    value. Values already at logical size pass through; a bare float passes through for the
    gate to validate (accepted only when `num_rows == 1`).
    """
    if isinstance(values, (int, float)):
        return values
    rows = self.num_logical_rows
    t = torch.as_tensor(values).squeeze()
    if t.ndim > 1:
        raise ValueError(
            f"Gate readout returned a tensor of shape {tuple(torch.as_tensor(values).shape)}; "
            f"expected per-row values of shape [B] (extra dimensions must be size 1)."
        )
    flat = t.reshape(-1)
    if flat.numel() == rows:
        return flat
    if flat.numel() == hidden_batch and rows and hidden_batch % rows == 0:
        factor = hidden_batch // rows
        return flat[::factor]
    raise ValueError(
        f"Gate readout returned {flat.numel()} values for a hidden batch of "
        f"{hidden_batch} and {rows} logical row(s); return one value per hidden row or per "
        f"logical row."
    )

_extract_pass_positions(forward_kwargs, seq_len) staticmethod

The pass's cache-axis positions from the hooked module's kwargs, when recoverable.

Prefers the cache_position kwarg (transformers threads it into decoder layers only when the model caller passes it explicitly). When absent, the padded-cache-axis offset is reconstructed from position_ids, which transformers v5 forwards to decoder layers on every pass: the longest row of a batch is unpadded, so its positions coincide with the padded cache axis and max(position_ids) - (seq_len - 1) is the pass's start offset. Neither kwarg reaches non-decoder hook points (e.g. o_proj inputs), which keep the counting fallback.

Source code in steerability/algorithms/state_control/common/runtime.py

@staticmethod
def _extract_pass_positions(forward_kwargs: dict | None, seq_len: int) -> torch.Tensor | None:
    """The pass's cache-axis positions from the hooked module's kwargs, when recoverable.

    Prefers the `cache_position` kwarg (transformers threads it into decoder layers only when
    the model caller passes it explicitly). When absent, the padded-cache-axis offset is
    reconstructed from `position_ids`, which transformers v5 forwards to decoder layers on
    every pass: the longest row of a batch is unpadded, so its positions coincide with the
    padded cache axis and `max(position_ids) - (seq_len - 1)` is the pass's start offset.
    Neither kwarg reaches non-decoder hook points (e.g. `o_proj` inputs), which keep the
    counting fallback.
    """
    if not forward_kwargs:
        return None
    positions = forward_kwargs.get("cache_position")
    if positions is not None and torch.is_tensor(positions) and positions.numel() > 0:
        return positions
    position_ids = forward_kwargs.get("position_ids")
    if position_ids is not None and torch.is_tensor(position_ids) and position_ids.numel() > 0:
        offset = int(position_ids.max().item()) - (seq_len - 1)
        if offset >= 0:
            return torch.tensor([offset], device=position_ids.device)
    return None

_position_offset(seq_len, cache_position, is_pass_opener)

Resolve the absolute position offset for the current pass, or None to skip the pass.

Auxiliary passes (marked via auxiliary_pass()) never advance the fallback counter. An aligned auxiliary pass is positioned by the pass positions when the hooked module supplies them (from cache_position or position_ids) and skipped otherwise; a detached auxiliary pass is always skipped. Ordinary passes take their offset from the pass positions when present. The opener maintains the fallback counter on every ordinary pass, so hook points without either kwarg keep the one-pass-one-step accounting unchanged and an anomalous pass missing them degrades to counting.

Parameters:

Name	Type	Description	Default
seq_len	int	The sequence length seen by this hook on this call.	required
cache_position	Tensor | None	The pass positions (from cache_position or reconstructed from position_ids), when the hooked module supplies them.	required
is_pass_opener	bool	Whether this hook is the designated pass opener.	required

Returns:

Type	Description
int | None	The absolute position offset to use when building the token mask, or None when the
int | None	pass must be skipped.

Warns:

Type	Description
UserWarning	Once per generation for each of: an aligned auxiliary pass at a hook point without pass positions (its transform is skipped); a multi-token ordinary pass after prefill at such a hook point (a multi-call decode pattern that counting cannot place); pass positions disappearing after having been observed.

Source code in steerability/algorithms/state_control/common/runtime.py

def _position_offset(
    self, seq_len: int, cache_position: torch.Tensor | None, is_pass_opener: bool
) -> int | None:
    """Resolve the absolute position offset for the current pass, or None to skip the pass.

    Auxiliary passes (marked via `auxiliary_pass()`) never advance the fallback counter. An
    aligned auxiliary pass is positioned by the pass positions when the hooked module supplies
    them (from `cache_position` or `position_ids`) and skipped otherwise; a detached auxiliary
    pass is always skipped. Ordinary passes take their offset from the pass positions when
    present. The opener maintains the fallback counter on every ordinary pass, so hook points
    without either kwarg keep the one-pass-one-step accounting unchanged and an anomalous pass
    missing them degrades to counting.

    Args:
        seq_len: The sequence length seen by this hook on this call.
        cache_position: The pass positions (from `cache_position` or reconstructed from
            `position_ids`), when the hooked module supplies them.
        is_pass_opener: Whether this hook is the designated pass opener.

    Returns:
        The absolute position offset to use when building the token mask, or None when the
        pass must be skipped.

    Warns:
        UserWarning: Once per generation for each of: an aligned auxiliary pass at a hook point
            without pass positions (its transform is skipped); a multi-token ordinary pass
            after prefill at such a hook point (a multi-call decode pattern that counting cannot
            place); pass positions disappearing after having been observed.
    """
    aux = current_auxiliary_pass()
    if aux is not None:
        if not aux.aligned:
            return None
        if cache_position is not None:
            return int(cache_position[0])
        self._warn_once(
            "aux_without_cache_position",
            "Auxiliary same-model passes cannot be position-mapped at this hook point (the "
            "hooked module receives neither `cache_position` nor `position_ids`); their "
            "transforms are skipped. Hook decoder layers for exact composition with same-model "
            "output controls.",
        )
        return None

    if is_pass_opener:
        if self._prefill_seen:
            if seq_len > 1 and cache_position is None:
                self._warn_once(
                    "multi_call_without_cache_position",
                    "Multiple generate calls detected at a hook point that receives neither "
                    "`cache_position` nor `position_ids`; position scoping may be skewed for "
                    "this generation.",
                )
            self._pass_offset = self._offset
            self._offset += seq_len
        else:
            self._pass_offset = 0
            self._offset = seq_len
            self._prefill_seen = True

    if cache_position is not None:
        self._clock_seen = True
        return int(cache_position[0])
    if self._clock_seen:
        self._warn_once(
            "inconsistent_cache_position",
            "Pass positions were available on earlier passes but missing on this one; "
            "falling back to pass counting for this pass.",
        )
    return self._pass_offset

_prefill_prompt_mask(hidden, pass_offset)

The stored prompt mask aligned to the hidden batch, on the prefill pass only.

The first pass may be longer than the prompt when a teacher-forced continuation is appended (e.g. compute_logprobs forwards [prompt; ref] in one pass). The continuation columns are not prompt, so the mask is extended with False there, and condition aggregation then covers exactly the real prompt tokens, reproducing generation-time scoring. A first pass shorter than the stored mask indicates misuse and raises.

Source code in steerability/algorithms/state_control/common/runtime.py

def _prefill_prompt_mask(self, hidden: torch.Tensor, pass_offset: int) -> torch.Tensor | None:
    """The stored prompt mask aligned to the hidden batch, on the prefill pass only.

    The first pass may be longer than the prompt when a teacher-forced continuation is
    appended (e.g. `compute_logprobs` forwards `[prompt; ref]` in one pass). The continuation
    columns are not prompt, so the mask is extended with False there, and condition
    aggregation then covers exactly the real prompt tokens, reproducing generation-time
    scoring. A first pass shorter than the stored mask indicates misuse and raises.
    """
    if self._prompt_mask is None or pass_offset != 0:
        return None
    width = self._prompt_mask.size(1)
    seq_len = hidden.size(1)
    if width > seq_len:
        raise ValueError(
            f"Prompt mask length {width} exceeds prefill sequence length {seq_len}."
        )
    mask = self._prompt_mask
    if width < seq_len:  # teacher-forced continuation appended after the prompt
        pad = torch.zeros(mask.size(0), seq_len - width, dtype=torch.bool)
        mask = torch.cat([mask, pad], dim=1)
    mask = align_mask_to_batch(mask, hidden.size(0))
    return mask.to(hidden.device)

_row_mask_for(gate, hidden)

Per-row gate decision expanded to the hidden batch as a [B_hidden, 1] mask.

Returns None when every row is closed (caller short-circuits). align_mask_to_batch performs the beam expansion and validates divisibility.

Source code in steerability/algorithms/state_control/common/runtime.py

def _row_mask_for(self, gate: Gate, hidden: torch.Tensor) -> torch.BoolTensor | None:
    """Per-row gate decision expanded to the hidden batch as a `[B_hidden, 1]` mask.

    Returns None when every row is closed (caller short-circuits). `align_mask_to_batch`
    performs the beam expansion and validates divisibility.
    """
    open_rows = gate.open_rows()
    if not bool(open_rows.any()):
        return None
    row_mask = align_mask_to_batch(open_rows.unsqueeze(1), hidden.size(0))  # [B_hidden, 1]
    return row_mask.to(hidden.device)

_warn_once(key, message)

Emit message as a UserWarning at most once per generation (keyed by key).

Source code in steerability/algorithms/state_control/common/runtime.py

def _warn_once(self, key: str, message: str) -> None:
    """Emit `message` as a UserWarning at most once per generation (keyed by `key`)."""
    if key in self._warned:
        return
    self._warned.add(key)
    warnings.warn(message, UserWarning)

build_behavior_hook(*, layer_id, transform, gate, token_scope, last_k=None, from_position=None, is_pass_opener=False, hook_point=None)

Build a hook that applies transform to the residual stream at layer_id, gated by gate.

The transform fires at the intersection of the token-scope mask and the gate's per-row decision (expanded across beams); a fully closed gate is a no-op, and a None gate leaves every row open.

Parameters:

Name	Type	Description	Default
layer_id	int	Index of the hooked layer (used to index per-layer transform artifacts).	required
transform	BaseTransform	The transform to apply at masked positions of open rows.	required
gate	Gate | None	Gate consulted per call, or None for unconditional application; row r of the hidden batch fires only when the gate's logical row r // beam_factor is open.	required
token_scope	ScopeKind	Which positions to steer (see make_token_mask).	required
last_k	int | None	Required when token_scope == "last_k".	None
from_position	int | None	Required when token_scope == "from_position".	None
is_pass_opener	bool	Whether this hook advances the shared position offset.	False
hook_point	HookPoint | None	Per-hook boundary override; defaults to the runtime's constructor value, so one runtime can host hooks at both boundaries.	None

Returns:

Type	Description
Callable	A hook callable suitable for the effective hook point (a forward hook for
Callable	"layer_output", a forward pre-hook for "layer_input").

Source code in steerability/algorithms/state_control/common/runtime.py

def build_behavior_hook(
    self,
    *,
    layer_id: int,
    transform: BaseTransform,
    gate: Gate | None,
    token_scope: ScopeKind,
    last_k: int | None = None,
    from_position: int | None = None,
    is_pass_opener: bool = False,
    hook_point: HookPoint | None = None,
) -> Callable:
    """Build a hook that applies `transform` to the residual stream at `layer_id`, gated by `gate`.

    The transform fires at the intersection of the token-scope mask and the gate's per-row
    decision (expanded across beams); a fully closed gate is a no-op, and a None gate
    leaves every row open.

    Args:
        layer_id: Index of the hooked layer (used to index per-layer transform artifacts).
        transform: The transform to apply at masked positions of open rows.
        gate: Gate consulted per call, or None for unconditional application; row `r` of
            the hidden batch fires only when the gate's logical row `r // beam_factor` is
            open.
        token_scope: Which positions to steer (see `make_token_mask`).
        last_k: Required when `token_scope == "last_k"`.
        from_position: Required when `token_scope == "from_position"`.
        is_pass_opener: Whether this hook advances the shared position offset.
        hook_point: Per-hook boundary override; defaults to the runtime's constructor
            value, so one runtime can host hooks at both boundaries.

    Returns:
        A hook callable suitable for the effective hook point (a forward hook for
        ``"layer_output"``, a forward pre-hook for ``"layer_input"``).
    """
    self._claim_opener(is_pass_opener)
    if (hook_point or self.hook_point) == "layer_output":

        def _forward_hook(module, args, kwargs, output):
            hidden = output[0] if isinstance(output, tuple) else output
            if hidden is None:
                return output
            hidden = self._apply(hidden, layer_id, transform, gate, token_scope, last_k,
                                 from_position, is_pass_opener, forward_kwargs=kwargs)
            return (hidden,) + output[1:] if isinstance(output, tuple) else hidden

        return _forward_hook

    def _pre_hook(module, input_args, input_kwargs):
        hidden = extract_hidden_states(input_args, input_kwargs)
        if hidden is None:
            return input_args, input_kwargs
        hidden = self._apply(hidden, layer_id, transform, gate, token_scope, last_k,
                             from_position, is_pass_opener, forward_kwargs=input_kwargs)
        return replace_hidden_states(input_args, input_kwargs, hidden)

    return _pre_hook

build_condition_hook(*, layer_id, gate, is_pass_opener=False, hook_point=None)

Build a read-only hook that scores the residual stream at layer_id and updates gate.

The hook never modifies hidden states. On each pass where the gate has not frozen its decision (not gate.is_ready()), it pools the hidden states per the gate's evidence (passing the pad-aware prompt mask on the prefill pass), computes per-row values via the evidence readout, collapses beam-expanded values to logical rows, and calls gate.update(rows, key=layer_id). Once the gate is ready, scoring is skipped entirely, and a gate whose rule never reports complete keeps re-scoring every pass. The hook still participates in pass-opener bookkeeping when the lowest hooked layer is a condition layer. Auxiliary passes are ignored entirely: no scoring, no gate update, no accounting.

Parameters:

Name	Type	Description	Default
layer_id	int	Index of the hooked layer.	required
gate	Gate	Gate whose evidence is read at this layer and fed the per-row values.	required
is_pass_opener	bool	Whether this hook advances the shared position offset.	False
hook_point	HookPoint | None	Per-hook boundary override; defaults to the runtime's constructor value.	None

Returns:

Type	Description
Callable	A hook callable suitable for the effective hook point.

Source code in steerability/algorithms/state_control/common/runtime.py

def build_condition_hook(
    self,
    *,
    layer_id: int,
    gate: Gate,
    is_pass_opener: bool = False,
    hook_point: HookPoint | None = None,
) -> Callable:
    """Build a read-only hook that scores the residual stream at `layer_id` and updates `gate`.

    The hook never modifies hidden states. On each pass where the gate has not frozen its
    decision (`not gate.is_ready()`), it pools the hidden states per the gate's evidence
    (passing the pad-aware prompt mask on the prefill pass), computes per-row values via
    the evidence readout, collapses beam-expanded values to logical rows, and calls
    `gate.update(rows, key=layer_id)`. Once the gate is ready, scoring is skipped
    entirely, and a gate whose rule never reports complete keeps re-scoring every pass.
    The hook still participates in pass-opener bookkeeping when the lowest hooked layer is
    a condition layer. Auxiliary passes are ignored entirely: no scoring, no gate update,
    no accounting.

    Args:
        layer_id: Index of the hooked layer.
        gate: Gate whose evidence is read at this layer and fed the per-row values.
        is_pass_opener: Whether this hook advances the shared position offset.
        hook_point: Per-hook boundary override; defaults to the runtime's constructor
            value.

    Returns:
        A hook callable suitable for the effective hook point.
    """
    self._claim_opener(is_pass_opener)

    def _score(hidden: torch.Tensor, forward_kwargs: dict | None) -> None:
        if current_auxiliary_pass() is not None:
            return
        cache_position = self._extract_pass_positions(forward_kwargs, hidden.size(1))
        pass_offset = self._position_offset(hidden.size(1), cache_position, is_pass_opener)
        if gate.is_ready():
            return
        prompt_mask = self._prefill_prompt_mask(hidden, pass_offset)
        pooled = aggregate_condition_hidden(
            hidden, gate.evidence.pooling, attention_mask=prompt_mask
        )
        values = gate.evidence.readout(pooled, layer_id)
        gate.update(self._collapse_to_rows(values, hidden.size(0)), key=layer_id)

    if (hook_point or self.hook_point) == "layer_output":

        def _forward_hook(module, args, kwargs, output):
            hidden = output[0] if isinstance(output, tuple) else output
            if hidden is None:
                return output
            _score(hidden, kwargs)
            return output

        return _forward_hook

    def _pre_hook(module, input_args, input_kwargs):
        hidden = extract_hidden_states(input_args, input_kwargs)
        if hidden is None:
            return input_args, input_kwargs
        _score(hidden, input_kwargs)
        return input_args, input_kwargs

    return _pre_hook

reset(prompt_lens, prompt_mask=None)

Clear position/prefill state and store the prompt lengths/mask for this generation.

Parameters:

Name	Type	Description	Default
prompt_lens	LongTensor	Per-row prompt lengths of shape [B_logical] (from compute_prompt_lens). Defines the logical batch size for row gating.	required
prompt_mask	Tensor | None	Optional pad-aware prompt attention mask of shape [B_logical, T_prompt] (True/1 at real tokens). Forwarded to condition scorers on the prefill pass so condition scores align with the real (non-pad) prompt positions, matching the mask the selector calibrated on.	None

Source code in steerability/algorithms/state_control/common/runtime.py

def reset(
    self,
    prompt_lens: torch.LongTensor,
    prompt_mask: torch.Tensor | None = None,
) -> None:
    """Clear position/prefill state and store the prompt lengths/mask for this generation.

    Args:
        prompt_lens: Per-row prompt lengths of shape ``[B_logical]`` (from
            `compute_prompt_lens`). Defines the logical batch size for row gating.
        prompt_mask: Optional pad-aware prompt attention mask of shape
            ``[B_logical, T_prompt]`` (True/1 at real tokens). Forwarded to condition scorers
            on the prefill pass so condition scores align with the real (non-pad) prompt
            positions, matching the mask the selector calibrated on.
    """
    self._prompt_lens = prompt_lens
    if prompt_mask is not None:
        mask = torch.as_tensor(prompt_mask)
        if mask.ndim == 1:
            mask = mask.unsqueeze(0)
        if mask.size(0) != prompt_lens.size(0):
            raise ValueError(
                f"prompt_mask has {mask.size(0)} rows but prompt_lens has "
                f"{prompt_lens.size(0)}; these must describe the same logical batch."
            )
        self._prompt_mask = mask.bool()
    else:
        self._prompt_mask = None
    self._offset = 0
    self._pass_offset = 0
    self._prefill_seen = False
    self._opener_built = False
    self._clock_seen = False
    self._warned = set()

_submodule_exists(model, module_path)

True when module_path resolves on the model's module tree.

Source code in steerability/algorithms/state_control/common/runtime.py

def _submodule_exists(model, module_path: str) -> bool:
    """True when `module_path` resolves on the model's module tree."""
    try:
        model.get_submodule(module_path)
    except AttributeError:
        return False
    return True

build_hooks(interventions, layout, prompt_lens, prompt_mask=None, model=None)

Compile bound interventions to torch hooks for one logical generation.

Creates a fresh TransformHookRuntime (per-generation position state is born here), resets every intervention's gate to the logical batch size (gate reset is idempotent, so a gate instance shared across interventions is reset harmlessly more than once), and emits one behavior hook per (intervention, layer) plus one condition hook per (intervention.gate.evidence, layer). An intervention whose gate is None builds no condition hooks and its behavior hooks apply with every row open; an intervention marked gate_driven_externally builds no condition hooks either, since another intervention's hooks feed its shared gate. Condition hooks precede behavior hooks so a gate update runs before the transform at a shared layer. Exactly one hook opens each pass: the first-firing hook of the lowest hooked layer across the tuple.

Module paths derive from each intervention's resolved site: decoder layers for residual transforms, the attention output projection for head_additive, and each layer's normalization sub-modules for the "norm_input" site. The intervention's boundary picks the hook phase ("layer_output" builds forward hooks, "layer_input" forward pre-hooks); the o_proj and "norm_input" sites hook module inputs.

Parameters:

Name	Type	Description	Default
interventions	Sequence['Intervention']	Bound interventions, in application order.	required
layout	'ModelFacts'	The module-path ModelLayout naming decoder layers, output projections, and norm sub-modules.	required
prompt_lens	Tensor	Per-row prompt lengths of shape [B_logical] (from compute_prompt_lens); defines the logical batch size for row gating.	required
prompt_mask	Tensor | None	Optional pad-aware prompt attention mask of shape [B_logical, T_prompt], forwarded to evidence pooling on the prefill pass.	None
model	'PreTrainedModel | None'	Optional live model, consulted only to skip norm sub-modules a layer does not define at the "norm_input" site.	None

Returns:

Type	Description
dict[str, list]	Hook specifications keyed by phase ("pre", "forward", "backward"), each entry a
dict[str, list]	mapping with "module" and "hook_func".

Raises:

Type	Description
ValueError	If an intervention is unbound, a layer has no module path in layout, or an o_proj-site intervention targets a layer that carries no attention module.

Source code in steerability/algorithms/state_control/common/runtime.py

def build_hooks(
    interventions: Sequence["Intervention"],
    layout: "ModelFacts",
    prompt_lens: torch.Tensor,
    prompt_mask: torch.Tensor | None = None,
    model: "PreTrainedModel | None" = None,
) -> dict[str, list]:
    """Compile bound interventions to torch hooks for one logical generation.

    Creates a fresh `TransformHookRuntime` (per-generation position state is born here), resets
    every intervention's gate to the logical batch size (gate reset is idempotent, so a gate
    instance shared across interventions is reset harmlessly more than once), and emits one
    behavior hook per (intervention, layer) plus one condition hook per
    (intervention.gate.evidence, layer). An intervention whose gate is None builds no condition
    hooks and its behavior hooks apply with every row open; an intervention marked
    `gate_driven_externally` builds no condition hooks either, since another intervention's
    hooks feed its shared gate. Condition hooks precede behavior hooks so a gate update runs
    before the transform at a shared layer. Exactly one hook opens each pass: the first-firing
    hook of the lowest hooked layer across the tuple.

    Module paths derive from each intervention's resolved site: decoder layers for residual
    transforms, the attention output projection for `head_additive`, and each layer's
    normalization sub-modules for the `"norm_input"` site. The intervention's `boundary` picks
    the hook phase (`"layer_output"` builds forward hooks, `"layer_input"` forward pre-hooks);
    the `o_proj` and `"norm_input"` sites hook module inputs.

    Args:
        interventions: Bound interventions, in application order.
        layout: The module-path `ModelLayout` naming decoder layers, output projections, and
            norm sub-modules.
        prompt_lens: Per-row prompt lengths of shape `[B_logical]` (from
            `compute_prompt_lens`); defines the logical batch size for row gating.
        prompt_mask: Optional pad-aware prompt attention mask of shape `[B_logical, T_prompt]`,
            forwarded to evidence pooling on the prefill pass.
        model: Optional live model, consulted only to skip norm sub-modules a layer does not
            define at the `"norm_input"` site.

    Returns:
        Hook specifications keyed by phase (`"pre"`, `"forward"`, `"backward"`), each entry a
        mapping with `"module"` and `"hook_func"`.

    Raises:
        ValueError: If an intervention is unbound, a layer has no module path in `layout`, or an
            o_proj-site intervention targets a layer that carries no attention module.
    """
    from .specs import Intervention

    runtime = TransformHookRuntime()
    runtime.reset(prompt_lens, prompt_mask)
    num_rows = int(prompt_lens.size(0))

    # hook units in module firing order: (layer, site_rank, condition_before_behavior)
    site_rank = {"pre": 0, "norm": 1, "o_proj": 2, "forward": 3}
    units: list[tuple[tuple, dict]] = []

    for intervention in interventions:
        if not isinstance(intervention, Intervention) or not isinstance(intervention.layers, tuple):
            raise ValueError("build_hooks requires bound interventions; call bind() first.")
        gate = intervention.gate
        if gate is not None and not isinstance(gate, Gate):
            raise ValueError("build_hooks requires a resolved gate; call bind() first.")
        if gate is not None:
            gate.reset(num_rows)

        site = intervention.resolved_site()
        boundary = intervention.boundary

        if gate is not None and not intervention.gate_driven_externally:
            for layer_id in gate.evidence.layer_ids:
                phase = "forward" if boundary == "layer_output" else "pre"
                units.append((
                    (layer_id, site_rank[phase if phase == "pre" else "forward"], 0),
                    {
                        "kind": "condition", "phase": phase, "layer_id": layer_id,
                        "module": layout.layer_names[layer_id],
                        "gate": gate, "hook_point": boundary,
                    },
                ))

        for layer_id in intervention.layers:
            if site == "norm_input":
                for norm_attr in layout.norm_attrs:
                    module = f"{layout.layer_names[layer_id]}.{norm_attr}"
                    if model is not None and not _submodule_exists(model, module):
                        continue
                    units.append((
                        (layer_id, site_rank["norm"], 1, module),
                        {
                            "kind": "behavior", "phase": "pre", "layer_id": layer_id,
                            "module": module, "intervention": intervention, "gate": gate,
                            "hook_point": "layer_input",
                        },
                    ))
            elif site == "o_proj":
                if not layout.has_attention(layer_id):
                    raise ValueError(
                        f"Intervention at the o_proj site targets decoder layer {layer_id}, which "
                        f"carries no attention module; attention layers of this model are "
                        f"{list(layout.attention_layers)}. Restrict the behavior layers to those, "
                        "or use a residual-stream site (decoder_layer or norm_input)."
                    )
                units.append((
                    (layer_id, site_rank["o_proj"], 1),
                    {
                        "kind": "behavior", "phase": "pre", "layer_id": layer_id,
                        "module": layout.oproj_names[layer_id], "intervention": intervention,
                        "gate": gate, "hook_point": "layer_input",
                    },
                ))
            else:
                phase = "forward" if boundary == "layer_output" else "pre"
                units.append((
                    (layer_id, site_rank[phase], 1),
                    {
                        "kind": "behavior", "phase": phase, "layer_id": layer_id,
                        "module": layout.layer_names[layer_id], "intervention": intervention,
                        "gate": gate, "hook_point": boundary,
                    },
                ))

    hooks: dict[str, list] = {"pre": [], "forward": [], "backward": []}
    if not units:
        return hooks

    # exactly one opener: the first-registered unit among those sharing the minimal firing key
    opener_index = min(range(len(units)), key=lambda index: units[index][0])

    for index, (key, unit) in enumerate(units):
        is_opener = index == opener_index
        if unit["kind"] == "condition":
            hook_func = runtime.build_condition_hook(
                layer_id=unit["layer_id"],
                gate=unit["gate"],
                is_pass_opener=is_opener,
                hook_point=unit["hook_point"],
            )
        else:
            intervention = unit["intervention"]
            hook_func = runtime.build_behavior_hook(
                layer_id=unit["layer_id"],
                transform=intervention.transform,
                gate=unit["gate"],
                token_scope=intervention.scope.kind,
                last_k=intervention.scope.last_k,
                from_position=intervention.scope.from_position,
                is_pass_opener=is_opener,
                hook_point=unit["hook_point"],
            )
        hooks[unit["phase"]].append({"module": unit["module"], "hook_func": hook_func})

    return hooks

selectors

Selector components for state control.

base

Base class for selectors that choose control sites and parameters.

T = TypeVar('T') module-attribute

BaseSelector

Bases: ABC, Generic[T]

Selects where and how to apply a control, given a model and artifacts.

The type parameter T describes the selection output (e.g., a list of layer ids, or a named tuple of (layer_id, threshold, comparator)).

Source code in steerability/algorithms/state_control/common/selectors/base.py

class BaseSelector(ABC, Generic[T]):
    """Selects where and how to apply a control, given a model and artifacts.

    The type parameter T describes the selection output (e.g., a list of layer
    ids, or a named tuple of (layer_id, threshold, comparator)).
    """

    @abstractmethod
    def select(self, **kwargs) -> T:
        """Run selection logic and return the result.

        Returns:
            Selection result of type T.
        """
        ...

select(**kwargs) abstractmethod

Run selection logic and return the result.

Returns:

Type	Description
T	Selection result of type T.

Source code in steerability/algorithms/state_control/common/selectors/base.py

@abstractmethod
def select(self, **kwargs) -> T:
    """Run selection logic and return the result.

    Returns:
        Selection result of type T.
    """
    ...

condition_point

Condition point search: find optimal (layer, threshold, comparator).

logger = logging.getLogger(__name__) module-attribute

ConditionPoint dataclass

Result of a condition point search, reusable as a self-contained CAST condition argument.

A ConditionPoint carries everything CAST needs to gate without re-searching: the layer, threshold, comparator, and (when populated by the selector) the runtime aggregation mode. Pass one directly to CASTArgs.condition_point to reuse a searched point, or flipped() to condition on the complement.

Attributes:

Name	Type	Description
layer_id	int	The condition layer (0-based).
threshold	float	The gate threshold.
comparator	Comparator	The gate comparator ("ge" opens when score >= threshold; "le" when score <= threshold).
f1	float	F1 of the class separation achieved by the search at this point.
margin	float	Geometric margin: distance from the decision boundary to the nearest calibration point. Positive means the calibration classes are cleanly separated at this threshold.
comparison_mode	CompMode | None	The runtime token-aggregation mode ("mean" or "last") the search was run for, or None when unset. Populated by ConditionPointSelector.select.

Source code in steerability/algorithms/state_control/common/selectors/condition_point.py

@dataclass
class ConditionPoint:
    """Result of a condition point search, reusable as a self-contained CAST condition argument.

    A `ConditionPoint` carries everything CAST needs to gate without re-searching: the layer,
    threshold, comparator, and (when populated by the selector) the runtime aggregation mode. Pass
    one directly to `CASTArgs.condition_point` to reuse a searched point, or `flipped()` to condition
    on the complement.

    Attributes:
        layer_id: The condition layer (0-based).
        threshold: The gate threshold.
        comparator: The gate comparator ("ge" opens when score >= threshold; "le" when
            score <= threshold).
        f1: F1 of the class separation achieved by the search at this point.
        margin: Geometric margin: distance from the decision boundary to the nearest calibration
            point. Positive means the calibration classes are cleanly separated at this threshold.
        comparison_mode: The runtime token-aggregation mode ("mean" or "last") the search was run
            for, or None when unset. Populated by `ConditionPointSelector.select`.
    """

    layer_id: int
    threshold: float
    comparator: Comparator
    f1: float
    margin: float = 0.0
    comparison_mode: CompMode | None = None

    def flipped(self) -> "ConditionPoint":
        """Return a copy with the comparator inverted ("ge" <-> "le").

        Conditions on the complement of this point (e.g. gate on scores at or below the threshold
        instead of at or above it). All other fields are preserved: `f1` and `margin` describe the
        original search and are carried over unchanged.
        """
        flipped_comparator: Comparator = "le" if self.comparator == "ge" else "ge"
        return ConditionPoint(
            layer_id=self.layer_id,
            threshold=self.threshold,
            comparator=flipped_comparator,
            f1=self.f1,
            margin=self.margin,
            comparison_mode=self.comparison_mode,
        )

comparator instance-attribute

comparison_mode = None class-attribute instance-attribute

f1 instance-attribute

layer_id instance-attribute

margin = 0.0 class-attribute instance-attribute

threshold instance-attribute

__init__(layer_id, threshold, comparator, f1, margin=0.0, comparison_mode=None)

flipped()

Return a copy with the comparator inverted ("ge" <-> "le").

Conditions on the complement of this point (e.g. gate on scores at or below the threshold instead of at or above it). All other fields are preserved: f1 and margin describe the original search and are carried over unchanged.

Source code in steerability/algorithms/state_control/common/selectors/condition_point.py

def flipped(self) -> "ConditionPoint":
    """Return a copy with the comparator inverted ("ge" <-> "le").

    Conditions on the complement of this point (e.g. gate on scores at or below the threshold
    instead of at or above it). All other fields are preserved: `f1` and `margin` describe the
    original search and are carried over unchanged.
    """
    flipped_comparator: Comparator = "le" if self.comparator == "ge" else "ge"
    return ConditionPoint(
        layer_id=self.layer_id,
        threshold=self.threshold,
        comparator=flipped_comparator,
        f1=self.f1,
        margin=self.margin,
        comparison_mode=self.comparison_mode,
    )

ConditionPointSelector

Bases: BaseSelector[ConditionPoint]

Grid-searches for the (layer, threshold, comparator) that best separates positive from negative examples using projected cosine similarity.

Examples are rendered via render_for_model according to fit_spec.prompt_format and tokenized with add_special_tokens=False for chat-templated text (matching the inference rendering of the condition gate).

The returned comparator is one of "ge"/"le" ("ge" opens the gate when score >= threshold), consumed directly by the PerKeyThreshold gate rule.

Source code in steerability/algorithms/state_control/common/selectors/condition_point.py

class ConditionPointSelector(BaseSelector[ConditionPoint]):
    """Grid-searches for the (layer, threshold, comparator) that best
    separates positive from negative examples using projected cosine
    similarity.

    Examples are rendered via `render_for_model` according to
    `fit_spec.prompt_format` and tokenized with `add_special_tokens=False` for
    chat-templated text (matching the inference rendering of the condition gate).

    The returned `comparator` is one of "ge"/"le" ("ge" opens the gate when
    score >= threshold), consumed directly by the `PerKeyThreshold` gate rule.
    """

    def select(
        self,
        *,
        model: PreTrainedModel | None,
        tokenizer: PreTrainedTokenizerBase,
        condition_directions: dict[int, torch.Tensor],
        data: ContrastivePairs,
        fit_spec: VectorTrainSpec,
        search_spec: ConditionSearchSpec,
        comparison_mode: CompMode | None = None,
        score: Literal["projected_cosine", "cosine"] = "projected_cosine",
        session: "SteeringSession | None" = None,
    ) -> ConditionPoint:
        """Run the grid search.

        Calibration extracts hidden states at `location="layer_input"`, matching the residual-stream
        boundary CAST's runtime condition pre-hook observes, and scores with the same score function
        the runtime scorer will apply, so selector and runtime scores agree for the same input.
        A `fit_spec` whose `location` is not `"layer_input"` means the condition direction was fit
        at a different boundary than it is calibrated and scored against, and triggers a
        `UserWarning`.

        Args:
            model: Model for hidden state extraction.
            tokenizer: Tokenizer for encoding.
            condition_directions: Per-layer direction tensors (from estimator).
            data: Contrastive pairs to evaluate separation on.
            fit_spec: How hidden states were accumulated (for span selection).
            search_spec: Search grid configuration.
            comparison_mode: The runtime condition aggregation mode CAST will use. Accepted for
                caller symmetry with CAST; calibration pools over `fit_spec.accumulate` spans.
            session: A `SteeringSession` serving hidden-state capture when no live model is
                available.
            score: Score function applied to the calibration examples, matching the runtime
                scorer the caller will build. `"projected_cosine"` scores a state as the cosine
                similarity with its tanh'd rank-one projection onto the direction, which is
                approximately `|cos(h, d)|` and therefore unsigned. `"cosine"` scores the signed
                cosine similarity against the direction; signed scores can be negative, so pair
                it with a `search_spec.threshold_range` that admits negative values, e.g.
                `(-1.0, 1.0)`.

        Returns:
            ConditionPoint with the best (layer, threshold, comparator, f1).

        Raises:
            ValueError: If `score` is not `"projected_cosine"` or `"cosine"`.
        """
        if score not in ("projected_cosine", "cosine"):
            raise ValueError(
                f"score must be 'projected_cosine' or 'cosine', got {score!r}."
            )

        device = next(model.parameters()).device if model is not None else torch.device("cpu")

        if fit_spec.location != "layer_input":
            warnings.warn(
                "ConditionPointSelector calibrates on layer-input hidden states (the boundary CAST's "
                f"runtime condition pre-hook observes), but fit_spec.location={fit_spec.location!r}. "
                "The condition direction was fit at a different residual-stream boundary than it is "
                "calibrated and scored against. Fit the condition vector with location='layer_input' "
                "to align fit, calibration, and runtime scoring.",
                UserWarning,
            )

        # render full texts according to prompt_format (shared with inference)
        rendered = render_contrastive(tokenizer, data, fit_spec.prompt_format)

        # tokenize
        enc_pos = tokenize_texts(tokenizer, rendered.pos_texts, device, add_special_tokens=rendered.add_special_tokens)
        enc_neg = tokenize_texts(tokenizer, rendered.neg_texts, device, add_special_tokens=rendered.add_special_tokens)

        # extract hidden states at the layer-input boundary the runtime pre-hook observes
        hs_pos, mask_pos = capture_hidden(
            enc_pos, model=model, session=session, batch_size=fit_spec.batch_size, location="layer_input"
        )
        hs_neg, mask_neg = capture_hidden(
            enc_neg, model=model, session=session, batch_size=fit_spec.batch_size, location="layer_input"
        )

        # spans are mask-derived, so the returned mask keeps them aligned when a remote
        # capture re-packs rows
        def _span_enc(enc, mask):
            if mask is None:
                return {k: v.cpu() for k, v in enc.items()}
            return {"input_ids": mask, "attention_mask": mask}

        # tokenize prompts separately if needed
        prompt_enc = None
        if fit_spec.accumulate == "suffix-only" and rendered.prompt_texts is not None:
            prompt_enc = tokenize_texts(
                tokenizer, rendered.prompt_texts, device, add_special_tokens=rendered.add_special_tokens
            )
            prompt_enc = {k: v.cpu() for k, v in prompt_enc.items()}

        spans_pos = select_spans(_span_enc(enc_pos, mask_pos), prompt_enc, fit_spec.accumulate)
        spans_neg = select_spans(_span_enc(enc_neg, mask_neg), prompt_enc, fit_spec.accumulate)

        # determine layers to search (0-based, matching runtime condition layer ids)
        if search_spec.candidate_layers is not None:
            layers = list(search_spec.candidate_layers)
        else:
            start, end = search_spec.layer_range or (0, len(hs_pos))
            layers = list(range(start, min(end, len(hs_pos))))

        grid = _threshold_grid(search_spec.threshold_range, search_spec.threshold_step)

        best = {"f1": -1.0, "margin": float("-inf"), "layer": 0, "thr": 0.0, "direction": "ge"}

        logger.debug("Searching %d layers with %d threshold values", len(layers), len(grid))

        for lid in layers:
            if lid not in condition_directions:
                continue

            Hp = pool_over_spans(hs_pos[lid], spans_pos)
            Hn = pool_over_spans(hs_neg[lid], spans_neg)
            c = condition_directions[lid].to(device=Hp.device, dtype=Hp.dtype)
            # squeeze [K, D] → [D] for K=1 (unified SteeringVector format)
            if c.ndim == 2 and c.shape[0] == 1:
                c = c.squeeze(0)

            if score == "cosine":
                sims_p = F.cosine_similarity(Hp.float(), c.float().unsqueeze(0), dim=-1)
                sims_n = F.cosine_similarity(Hn.float(), c.float().unsqueeze(0), dim=-1)
            else:
                projector = rank_one_projector(c)
                sims_p = projected_cosine_similarity_tensor(Hp, projector)
                sims_n = projected_cosine_similarity_tensor(Hn, projector)

            cand = _best_point_for_layer(sims_p, sims_n, grid)
            if (cand["f1"], cand["margin"]) > (best["f1"], best["margin"]):
                best.update(
                    f1=cand["f1"],
                    margin=cand["margin"],
                    layer=lid,
                    thr=cand["thr"],
                    direction=cand["comparator"],
                )

        logger.debug(
            "Best condition point: layer=%d, threshold=%.3f, comparator=%s, f1=%.3f, margin=%.4f",
            best["layer"], best["thr"], best["direction"], best["f1"], best["margin"],
        )

        if fit_spec.method == "mean_diff" and best["direction"] == "le":
            warnings.warn(
                f"Condition search selected comparator 'le' at layer {best['layer']} "
                f"(margin {best['margin']:.4f}). The direction was fit as mean(positives) - "
                "mean(negatives), so positives are expected to score HIGHER. An inverted "
                "comparator usually means the calibration set is too small or this layer carries "
                "little signal for the condition. Consider a larger calibration set, a narrower "
                "layer_range, or pinning condition_layer_ids/condition_vector_threshold.",
                UserWarning,
                stacklevel=2,
            )

        return ConditionPoint(
            layer_id=best["layer"],
            threshold=best["thr"],
            comparator=best["direction"],
            f1=best["f1"],
            margin=best["margin"],
            comparison_mode=comparison_mode,
        )

select(*, model, tokenizer, condition_directions, data, fit_spec, search_spec, comparison_mode=None, score='projected_cosine', session=None)

Run the grid search.

Calibration extracts hidden states at location="layer_input", matching the residual-stream boundary CAST's runtime condition pre-hook observes, and scores with the same score function the runtime scorer will apply, so selector and runtime scores agree for the same input. A fit_spec whose location is not "layer_input" means the condition direction was fit at a different boundary than it is calibrated and scored against, and triggers a UserWarning.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	Model for hidden state extraction.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for encoding.	required
condition_directions	dict[int, Tensor]	Per-layer direction tensors (from estimator).	required
data	ContrastivePairs	Contrastive pairs to evaluate separation on.	required
fit_spec	VectorTrainSpec	How hidden states were accumulated (for span selection).	required
search_spec	ConditionSearchSpec	Search grid configuration.	required
comparison_mode	CompMode | None	The runtime condition aggregation mode CAST will use. Accepted for caller symmetry with CAST; calibration pools over fit_spec.accumulate spans.	None
session	'SteeringSession | None'	A SteeringSession serving hidden-state capture when no live model is available.	None
score	Literal['projected_cosine', 'cosine']	Score function applied to the calibration examples, matching the runtime scorer the caller will build. "projected_cosine" scores a state as the cosine similarity with its tanh'd rank-one projection onto the direction, which is approximately |cos(h, d)| and therefore unsigned. "cosine" scores the signed cosine similarity against the direction; signed scores can be negative, so pair it with a search_spec.threshold_range that admits negative values, e.g. (-1.0, 1.0).	'projected_cosine'

Returns:

Type	Description
ConditionPoint	ConditionPoint with the best (layer, threshold, comparator, f1).

Raises:

Type	Description
ValueError	If score is not "projected_cosine" or "cosine".

Source code in steerability/algorithms/state_control/common/selectors/condition_point.py

def select(
    self,
    *,
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase,
    condition_directions: dict[int, torch.Tensor],
    data: ContrastivePairs,
    fit_spec: VectorTrainSpec,
    search_spec: ConditionSearchSpec,
    comparison_mode: CompMode | None = None,
    score: Literal["projected_cosine", "cosine"] = "projected_cosine",
    session: "SteeringSession | None" = None,
) -> ConditionPoint:
    """Run the grid search.

    Calibration extracts hidden states at `location="layer_input"`, matching the residual-stream
    boundary CAST's runtime condition pre-hook observes, and scores with the same score function
    the runtime scorer will apply, so selector and runtime scores agree for the same input.
    A `fit_spec` whose `location` is not `"layer_input"` means the condition direction was fit
    at a different boundary than it is calibrated and scored against, and triggers a
    `UserWarning`.

    Args:
        model: Model for hidden state extraction.
        tokenizer: Tokenizer for encoding.
        condition_directions: Per-layer direction tensors (from estimator).
        data: Contrastive pairs to evaluate separation on.
        fit_spec: How hidden states were accumulated (for span selection).
        search_spec: Search grid configuration.
        comparison_mode: The runtime condition aggregation mode CAST will use. Accepted for
            caller symmetry with CAST; calibration pools over `fit_spec.accumulate` spans.
        session: A `SteeringSession` serving hidden-state capture when no live model is
            available.
        score: Score function applied to the calibration examples, matching the runtime
            scorer the caller will build. `"projected_cosine"` scores a state as the cosine
            similarity with its tanh'd rank-one projection onto the direction, which is
            approximately `|cos(h, d)|` and therefore unsigned. `"cosine"` scores the signed
            cosine similarity against the direction; signed scores can be negative, so pair
            it with a `search_spec.threshold_range` that admits negative values, e.g.
            `(-1.0, 1.0)`.

    Returns:
        ConditionPoint with the best (layer, threshold, comparator, f1).

    Raises:
        ValueError: If `score` is not `"projected_cosine"` or `"cosine"`.
    """
    if score not in ("projected_cosine", "cosine"):
        raise ValueError(
            f"score must be 'projected_cosine' or 'cosine', got {score!r}."
        )

    device = next(model.parameters()).device if model is not None else torch.device("cpu")

    if fit_spec.location != "layer_input":
        warnings.warn(
            "ConditionPointSelector calibrates on layer-input hidden states (the boundary CAST's "
            f"runtime condition pre-hook observes), but fit_spec.location={fit_spec.location!r}. "
            "The condition direction was fit at a different residual-stream boundary than it is "
            "calibrated and scored against. Fit the condition vector with location='layer_input' "
            "to align fit, calibration, and runtime scoring.",
            UserWarning,
        )

    # render full texts according to prompt_format (shared with inference)
    rendered = render_contrastive(tokenizer, data, fit_spec.prompt_format)

    # tokenize
    enc_pos = tokenize_texts(tokenizer, rendered.pos_texts, device, add_special_tokens=rendered.add_special_tokens)
    enc_neg = tokenize_texts(tokenizer, rendered.neg_texts, device, add_special_tokens=rendered.add_special_tokens)

    # extract hidden states at the layer-input boundary the runtime pre-hook observes
    hs_pos, mask_pos = capture_hidden(
        enc_pos, model=model, session=session, batch_size=fit_spec.batch_size, location="layer_input"
    )
    hs_neg, mask_neg = capture_hidden(
        enc_neg, model=model, session=session, batch_size=fit_spec.batch_size, location="layer_input"
    )

    # spans are mask-derived, so the returned mask keeps them aligned when a remote
    # capture re-packs rows
    def _span_enc(enc, mask):
        if mask is None:
            return {k: v.cpu() for k, v in enc.items()}
        return {"input_ids": mask, "attention_mask": mask}

    # tokenize prompts separately if needed
    prompt_enc = None
    if fit_spec.accumulate == "suffix-only" and rendered.prompt_texts is not None:
        prompt_enc = tokenize_texts(
            tokenizer, rendered.prompt_texts, device, add_special_tokens=rendered.add_special_tokens
        )
        prompt_enc = {k: v.cpu() for k, v in prompt_enc.items()}

    spans_pos = select_spans(_span_enc(enc_pos, mask_pos), prompt_enc, fit_spec.accumulate)
    spans_neg = select_spans(_span_enc(enc_neg, mask_neg), prompt_enc, fit_spec.accumulate)

    # determine layers to search (0-based, matching runtime condition layer ids)
    if search_spec.candidate_layers is not None:
        layers = list(search_spec.candidate_layers)
    else:
        start, end = search_spec.layer_range or (0, len(hs_pos))
        layers = list(range(start, min(end, len(hs_pos))))

    grid = _threshold_grid(search_spec.threshold_range, search_spec.threshold_step)

    best = {"f1": -1.0, "margin": float("-inf"), "layer": 0, "thr": 0.0, "direction": "ge"}

    logger.debug("Searching %d layers with %d threshold values", len(layers), len(grid))

    for lid in layers:
        if lid not in condition_directions:
            continue

        Hp = pool_over_spans(hs_pos[lid], spans_pos)
        Hn = pool_over_spans(hs_neg[lid], spans_neg)
        c = condition_directions[lid].to(device=Hp.device, dtype=Hp.dtype)
        # squeeze [K, D] → [D] for K=1 (unified SteeringVector format)
        if c.ndim == 2 and c.shape[0] == 1:
            c = c.squeeze(0)

        if score == "cosine":
            sims_p = F.cosine_similarity(Hp.float(), c.float().unsqueeze(0), dim=-1)
            sims_n = F.cosine_similarity(Hn.float(), c.float().unsqueeze(0), dim=-1)
        else:
            projector = rank_one_projector(c)
            sims_p = projected_cosine_similarity_tensor(Hp, projector)
            sims_n = projected_cosine_similarity_tensor(Hn, projector)

        cand = _best_point_for_layer(sims_p, sims_n, grid)
        if (cand["f1"], cand["margin"]) > (best["f1"], best["margin"]):
            best.update(
                f1=cand["f1"],
                margin=cand["margin"],
                layer=lid,
                thr=cand["thr"],
                direction=cand["comparator"],
            )

    logger.debug(
        "Best condition point: layer=%d, threshold=%.3f, comparator=%s, f1=%.3f, margin=%.4f",
        best["layer"], best["thr"], best["direction"], best["f1"], best["margin"],
    )

    if fit_spec.method == "mean_diff" and best["direction"] == "le":
        warnings.warn(
            f"Condition search selected comparator 'le' at layer {best['layer']} "
            f"(margin {best['margin']:.4f}). The direction was fit as mean(positives) - "
            "mean(negatives), so positives are expected to score HIGHER. An inverted "
            "comparator usually means the calibration set is too small or this layer carries "
            "little signal for the condition. Consider a larger calibration set, a narrower "
            "layer_range, or pinning condition_layer_ids/condition_vector_threshold.",
            UserWarning,
            stacklevel=2,
        )

    return ConditionPoint(
        layer_id=best["layer"],
        threshold=best["thr"],
        comparator=best["direction"],
        f1=best["f1"],
        margin=best["margin"],
        comparison_mode=comparison_mode,
    )

_best_point_for_layer(sims_p, sims_n, grid)

Best (f1, margin, threshold, comparator) for one layer's calibration scores.

F1 saturates at 1.0 across many (threshold, comparator) points when the calibration set is small -- 10 pairs against a 100-point grid and 2 comparators leaves thousands of candidates, so a perfect fit carries little evidence. Ties are therefore broken by the geometric margin: the distance from the decision boundary to the nearest calibration point. Maximising it centres the threshold in the widest gap and prefers layers whose classes are genuinely far apart over layers that separate by luck.

Parameters:

Name	Type	Description	Default
sims_p	Tensor	[N_pos] condition scores for positives.	required
sims_n	Tensor	[N_neg] condition scores for negatives.	required
grid	Tensor	Candidate thresholds.	required

Returns:

Type	Description
dict	dict with keys "f1", "margin", "thr", "comparator".

Source code in steerability/algorithms/state_control/common/selectors/condition_point.py

def _best_point_for_layer(
    sims_p: torch.Tensor,
    sims_n: torch.Tensor,
    grid: torch.Tensor,
) -> dict:
    """Best (f1, margin, threshold, comparator) for one layer's calibration scores.

    F1 saturates at 1.0 across many (threshold, comparator) points when the calibration set is
    small -- 10 pairs against a 100-point grid and 2 comparators leaves thousands of candidates,
    so a perfect fit carries little evidence. Ties are therefore broken by the geometric margin:
    the distance from the decision boundary to the nearest calibration point. Maximising it
    centres the threshold in the widest gap and prefers layers whose classes are genuinely far
    apart over layers that separate by luck.

    Args:
        sims_p: [N_pos] condition scores for positives.
        sims_n: [N_neg] condition scores for negatives.
        grid: Candidate thresholds.

    Returns:
        dict with keys "f1", "margin", "thr", "comparator".
    """
    best = {"f1": -1.0, "margin": float("-inf"), "thr": 0.0, "comparator": "ge"}

    for cmp in ("ge", "le"):
        for thr in grid:
            thr_f = float(thr)

            if cmp == "ge":  # gate opens when score >= threshold
                tp = int((sims_p >= thr_f).sum().item())
                fp = int((sims_n >= thr_f).sum().item())
                margin = min(
                    float((sims_p - thr_f).min().item()),
                    float((thr_f - sims_n).min().item()),
                )
            else:  # gate opens when score <= threshold
                tp = int((sims_p <= thr_f).sum().item())
                fp = int((sims_n <= thr_f).sum().item())
                margin = min(
                    float((thr_f - sims_p).min().item()),
                    float((sims_n - thr_f).min().item()),
                )

            fn = int(sims_p.numel()) - tp
            prec = tp / (tp + fp + 1e-8)
            rec = tp / (tp + fn + 1e-8)
            f1 = 0.0 if prec + rec < 1e-8 else 2 * prec * rec / (prec + rec)

            if (f1, margin) > (best["f1"], best["margin"]):
                best = {"f1": f1, "margin": margin, "thr": thr_f, "comparator": cmp}

    return best

_threshold_grid(threshold_range, step)

Build a half-open, step-exact threshold grid [low, high).

Parameters:

Name	Type	Description	Default
threshold_range	tuple[float, float]	(low, high) bounds.	required
step	float	Grid increment (must be positive).	required

Returns:

Type	Description
Tensor	1-D tensor of candidate thresholds; [low] when the range admits no full step.

Source code in steerability/algorithms/state_control/common/selectors/condition_point.py

def _threshold_grid(threshold_range: tuple[float, float], step: float) -> torch.Tensor:
    """Build a half-open, step-exact threshold grid `[low, high)`.

    Args:
        threshold_range: `(low, high)` bounds.
        step: Grid increment (must be positive).

    Returns:
        1-D tensor of candidate thresholds; `[low]` when the range admits no full step.
    """
    low, high = threshold_range
    grid = torch.arange(low, high, step, dtype=torch.float64)
    if grid.numel() == 0:
        grid = torch.tensor([low], dtype=torch.float64)
    return grid

fixed_layer

Selector that returns an explicit, user-specified layer.

FixedLayerSelector

Bases: BaseSelector[int]

Selects a single layer specified at construction time.

Use when the caller already knows the target layer (e.g., from a config file or a previous search).

Parameters:

Name	Type	Description	Default
layer_id	int	The layer index to select.	required

Source code in steerability/algorithms/state_control/common/selectors/fixed_layer.py

class FixedLayerSelector(BaseSelector[int]):
    """Selects a single layer specified at construction time.

    Use when the caller already knows the target layer (e.g., from a config
    file or a previous search).

    Args:
        layer_id: The layer index to select.
    """

    component_kind: ClassVar[str] = "fixed_layer"

    def __init__(self, layer_id: int):
        if layer_id < 0:
            raise ValueError(f"layer_id must be >= 0, got {layer_id}.")
        self.layer_id = layer_id

    def to_config(self) -> dict:
        """The selector's serialized parameters."""
        return {"layer_id": int(self.layer_id)}

    @classmethod
    def from_config(cls, params: dict) -> "FixedLayerSelector":
        """Rebuild the selector from its serialized parameters."""
        return cls(layer_id=params["layer_id"])

    def select(self, *, num_layers: int) -> int:
        """Return the fixed layer id after bounds-checking.

        Args:
            num_layers: Total number of layers in the model.

        Returns:
            The stored layer id.

        Raises:
            ValueError: If layer_id >= num_layers.
        """
        if self.layer_id >= num_layers:
            raise ValueError(
                f"layer_id {self.layer_id} out of range [0, {num_layers})."
            )
        return self.layer_id

component_kind = 'fixed_layer' class-attribute

layer_id = layer_id instance-attribute

__init__(layer_id)

Source code in steerability/algorithms/state_control/common/selectors/fixed_layer.py

def __init__(self, layer_id: int):
    if layer_id < 0:
        raise ValueError(f"layer_id must be >= 0, got {layer_id}.")
    self.layer_id = layer_id

from_config(params) classmethod

Rebuild the selector from its serialized parameters.

Source code in steerability/algorithms/state_control/common/selectors/fixed_layer.py

@classmethod
def from_config(cls, params: dict) -> "FixedLayerSelector":
    """Rebuild the selector from its serialized parameters."""
    return cls(layer_id=params["layer_id"])

select(*, num_layers)

Return the fixed layer id after bounds-checking.

Parameters:

Name	Type	Description	Default
num_layers	int	Total number of layers in the model.	required

Returns:

Type	Description
int	The stored layer id.

Raises:

Type	Description
ValueError	If layer_id >= num_layers.

Source code in steerability/algorithms/state_control/common/selectors/fixed_layer.py

def select(self, *, num_layers: int) -> int:
    """Return the fixed layer id after bounds-checking.

    Args:
        num_layers: Total number of layers in the model.

    Returns:
        The stored layer id.

    Raises:
        ValueError: If layer_id >= num_layers.
    """
    if self.layer_id >= num_layers:
        raise ValueError(
            f"layer_id {self.layer_id} out of range [0, {num_layers})."
        )
    return self.layer_id

to_config()

The selector's serialized parameters.

Source code in steerability/algorithms/state_control/common/selectors/fixed_layer.py

def to_config(self) -> dict:
    """The selector's serialized parameters."""
    return {"layer_id": int(self.layer_id)}

fractional_depth

Selector that picks a layer at a given fractional depth of the model.

FractionalDepthSelector

Bases: BaseSelector[int]

Selects a layer at a specified fraction of total model depth.

For example, FractionalDepthSelector(0.4) on a 32-layer model returns layer 12 (≈ 40% depth), matching the CAA paper's finding that layer 13/32 is near-optimal.

Parameters:

Name	Type	Description	Default
fraction	float	Target depth as a fraction in (0, 1).	required
minimum	int	Floor for the returned layer id (default 0). Useful when very early layers should be excluded (e.g., ActAdd avoids layer 0).	0

Source code in steerability/algorithms/state_control/common/selectors/fractional_depth.py

class FractionalDepthSelector(BaseSelector[int]):
    """Selects a layer at a specified fraction of total model depth.

    For example, ``FractionalDepthSelector(0.4)`` on a 32-layer model returns
    layer 12 (≈ 40% depth), matching the CAA paper's finding that layer 13/32
    is near-optimal.

    Args:
        fraction: Target depth as a fraction in (0, 1).
        minimum: Floor for the returned layer id (default 0). Useful when
            very early layers should be excluded (e.g., ActAdd avoids layer 0).
    """

    component_kind: ClassVar[str] = "fractional_depth"

    def __init__(self, fraction: float, minimum: int = 0):
        if not 0.0 < fraction < 1.0:
            raise ValueError(f"fraction must be in (0, 1), got {fraction}.")
        if minimum < 0:
            raise ValueError(f"minimum must be >= 0, got {minimum}.")
        self.fraction = fraction
        self.minimum = minimum

    def to_config(self) -> dict:
        """The selector's serialized parameters."""
        return {"fraction": float(self.fraction), "minimum": int(self.minimum)}

    @classmethod
    def from_config(cls, params: dict) -> "FractionalDepthSelector":
        """Rebuild the selector from its serialized parameters."""
        return cls(fraction=params["fraction"], minimum=params.get("minimum", 0))

    def select(self, *, num_layers: int) -> int:
        """Compute and return the target layer id.

        Args:
            num_layers: Total number of layers in the model.

        Returns:
            Layer index at approximately ``fraction * num_layers``, clamped
            to ``[minimum, num_layers - 1]``.
        """
        layer_id = int(num_layers * self.fraction)
        layer_id = max(self.minimum, min(layer_id, num_layers - 1))
        return layer_id

component_kind = 'fractional_depth' class-attribute

fraction = fraction instance-attribute

minimum = minimum instance-attribute

__init__(fraction, minimum=0)

Source code in steerability/algorithms/state_control/common/selectors/fractional_depth.py

def __init__(self, fraction: float, minimum: int = 0):
    if not 0.0 < fraction < 1.0:
        raise ValueError(f"fraction must be in (0, 1), got {fraction}.")
    if minimum < 0:
        raise ValueError(f"minimum must be >= 0, got {minimum}.")
    self.fraction = fraction
    self.minimum = minimum

from_config(params) classmethod

Rebuild the selector from its serialized parameters.

Source code in steerability/algorithms/state_control/common/selectors/fractional_depth.py

@classmethod
def from_config(cls, params: dict) -> "FractionalDepthSelector":
    """Rebuild the selector from its serialized parameters."""
    return cls(fraction=params["fraction"], minimum=params.get("minimum", 0))

select(*, num_layers)

Compute and return the target layer id.

Parameters:

Name	Type	Description	Default
num_layers	int	Total number of layers in the model.	required

Returns:

Type	Description
int	Layer index at approximately fraction * num_layers, clamped
int	to [minimum, num_layers - 1].

Source code in steerability/algorithms/state_control/common/selectors/fractional_depth.py

def select(self, *, num_layers: int) -> int:
    """Compute and return the target layer id.

    Args:
        num_layers: Total number of layers in the model.

    Returns:
        Layer index at approximately ``fraction * num_layers``, clamped
        to ``[minimum, num_layers - 1]``.
    """
    layer_id = int(num_layers * self.fraction)
    layer_id = max(self.minimum, min(layer_id, num_layers - 1))
    return layer_id

to_config()

The selector's serialized parameters.

Source code in steerability/algorithms/state_control/common/selectors/fractional_depth.py

def to_config(self) -> dict:
    """The selector's serialized parameters."""
    return {"fraction": float(self.fraction), "minimum": int(self.minimum)}

LateThirdSelector

Bases: BaseSelector[list[int]]

Selects the late third of the model's layers.

The default behavior-layer heuristic for conditional activation steering.

Source code in steerability/algorithms/state_control/common/selectors/fractional_depth.py

class LateThirdSelector(BaseSelector[list[int]]):
    """Selects the late third of the model's layers.

    The default behavior-layer heuristic for conditional activation steering.
    """

    component_kind: ClassVar[str] = "late_third"

    def to_config(self) -> dict:
        """The selector's serialized parameters (none)."""
        return {}

    @classmethod
    def from_config(cls, params: dict) -> "LateThirdSelector":
        """Rebuild the selector from its serialized parameters."""
        return cls()

    def select(self, *, num_layers: int) -> list[int]:
        """Return the last third of the layer indices.

        Args:
            num_layers: Total number of layers in the model.

        Returns:
            The selected layer ids, ascending.
        """
        from .utils.layer_heuristics import late_third

        return late_third(num_layers)

component_kind = 'late_third' class-attribute

from_config(params) classmethod

Rebuild the selector from its serialized parameters.

Source code in steerability/algorithms/state_control/common/selectors/fractional_depth.py

@classmethod
def from_config(cls, params: dict) -> "LateThirdSelector":
    """Rebuild the selector from its serialized parameters."""
    return cls()

select(*, num_layers)

Return the last third of the layer indices.

Parameters:

Name	Type	Description	Default
num_layers	int	Total number of layers in the model.	required

Returns:

Type	Description
list[int]	The selected layer ids, ascending.

Source code in steerability/algorithms/state_control/common/selectors/fractional_depth.py

def select(self, *, num_layers: int) -> list[int]:
    """Return the last third of the layer indices.

    Args:
        num_layers: Total number of layers in the model.

    Returns:
        The selected layer ids, ascending.
    """
    from .utils.layer_heuristics import late_third

    return late_third(num_layers)

to_config()

The selector's serialized parameters (none).

Source code in steerability/algorithms/state_control/common/selectors/fractional_depth.py

def to_config(self) -> dict:
    """The selector's serialized parameters (none)."""
    return {}

top_k_head

Selector that returns top-K heads by probe accuracy.

logger = logging.getLogger(__name__) module-attribute

TopKHeadSelector

Bases: BaseSelector[list[tuple[int, int]]]

Selects the top-K (layer_id, head_id) pairs by probe accuracy.

Used by ITI to select which attention heads to intervene on based on linear probe training accuracy scores. "Probe" here is ITI's per-head fit-time classifier used to rank heads, distinct from the toolkit's Probe detector.

Parameters:

Name	Type	Description	Default
k	int	Number of top heads to select.	required

Source code in steerability/algorithms/state_control/common/selectors/top_k_head.py

class TopKHeadSelector(BaseSelector[list[tuple[int, int]]]):
    """Selects the top-K (layer_id, head_id) pairs by probe accuracy.

    Used by ITI to select which attention heads to intervene on based on
    linear probe training accuracy scores. "Probe" here is ITI's per-head fit-time
    classifier used to rank heads, distinct from the toolkit's `Probe` detector.

    Args:
        k: Number of top heads to select.
    """

    component_kind: ClassVar[str] = "top_k_head"

    def __init__(self, k: int):
        if k < 1:
            raise ValueError(f"k must be >= 1, got {k}.")
        self.k = k

    def to_config(self) -> dict:
        """The selector's serialized parameters."""
        return {"k": int(self.k)}

    @classmethod
    def from_config(cls, params: dict) -> "TopKHeadSelector":
        """Rebuild the selector from its serialized parameters."""
        return cls(k=params["k"])

    def select(self, *, steering_vector: SteeringVector) -> list[tuple[int, int]]:
        """Return the top-K (layer_id, head_id) pairs sorted by accuracy descending.

        Args:
            steering_vector: A SteeringVector with probe_accuracies populated.

        Returns:
            List of top-K (layer_id, head_id) tuples, sorted by accuracy descending.

        Raises:
            ValueError: If the steering vector has no probe_accuracies.
        """
        if steering_vector.probe_accuracies is None:
            raise ValueError("SteeringVector has no probe_accuracies for head selection.")
        ranked = sorted(steering_vector.probe_accuracies.items(), key=lambda item: item[1], reverse=True)
        selected = [key for key, _ in ranked[: self.k]]
        if len(selected) < self.k:
            logger.warning("Requested k=%d but only %d heads available.", self.k, len(selected))
        return selected

component_kind = 'top_k_head' class-attribute

k = k instance-attribute

__init__(k)

Source code in steerability/algorithms/state_control/common/selectors/top_k_head.py

def __init__(self, k: int):
    if k < 1:
        raise ValueError(f"k must be >= 1, got {k}.")
    self.k = k

from_config(params) classmethod

Rebuild the selector from its serialized parameters.

Source code in steerability/algorithms/state_control/common/selectors/top_k_head.py

@classmethod
def from_config(cls, params: dict) -> "TopKHeadSelector":
    """Rebuild the selector from its serialized parameters."""
    return cls(k=params["k"])

select(*, steering_vector)

Return the top-K (layer_id, head_id) pairs sorted by accuracy descending.

Parameters:

Name	Type	Description	Default
steering_vector	SteeringVector	A SteeringVector with probe_accuracies populated.	required

Returns:

Type	Description
list[tuple[int, int]]	List of top-K (layer_id, head_id) tuples, sorted by accuracy descending.

Raises:

Type	Description
ValueError	If the steering vector has no probe_accuracies.

Source code in steerability/algorithms/state_control/common/selectors/top_k_head.py

def select(self, *, steering_vector: SteeringVector) -> list[tuple[int, int]]:
    """Return the top-K (layer_id, head_id) pairs sorted by accuracy descending.

    Args:
        steering_vector: A SteeringVector with probe_accuracies populated.

    Returns:
        List of top-K (layer_id, head_id) tuples, sorted by accuracy descending.

    Raises:
        ValueError: If the steering vector has no probe_accuracies.
    """
    if steering_vector.probe_accuracies is None:
        raise ValueError("SteeringVector has no probe_accuracies for head selection.")
    ranked = sorted(steering_vector.probe_accuracies.items(), key=lambda item: item[1], reverse=True)
    selected = [key for key, _ in ranked[: self.k]]
    if len(selected) < self.k:
        logger.warning("Requested k=%d but only %d heads available.", self.k, len(selected))
    return selected

to_config()

The selector's serialized parameters.

Source code in steerability/algorithms/state_control/common/selectors/top_k_head.py

def to_config(self) -> dict:
    """The selector's serialized parameters."""
    return {"k": int(self.k)}

utils

Selector utility modules.

__all__ = ['late_third'] module-attribute

late_third(num_layers)

Return layer ids for the last third of the model.

This is the default heuristic for behavior layer selection when the user does not specify explicit layer ids.

Parameters:

Name	Type	Description	Default
num_layers	int	Total number of layers in the model.	required

Returns:

Type	Description
list[int]	List of layer ids.

Source code in steerability/algorithms/state_control/common/selectors/utils/layer_heuristics.py

def late_third(num_layers: int) -> list[int]:
    """Return layer ids for the last third of the model.

    This is the default heuristic for behavior layer selection when the
    user does not specify explicit layer ids.

    Args:
        num_layers: Total number of layers in the model.

    Returns:
        List of layer ids.
    """
    start = num_layers - (num_layers // 3)
    return list(range(start, num_layers))

layer_heuristics

Heuristic layer selection strategies.

late_third(num_layers)

Return layer ids for the last third of the model.

This is the default heuristic for behavior layer selection when the user does not specify explicit layer ids.

Parameters:

Name	Type	Description	Default
num_layers	int	Total number of layers in the model.	required

Returns:

Type	Description
list[int]	List of layer ids.

Source code in steerability/algorithms/state_control/common/selectors/utils/layer_heuristics.py

def late_third(num_layers: int) -> list[int]:
    """Return layer ids for the last third of the model.

    This is the default heuristic for behavior layer selection when the
    user does not specify explicit layer ids.

    Args:
        num_layers: Total number of layers in the model.

    Returns:
        List of layer ids.
    """
    start = num_layers - (num_layers // 3)
    return list(range(start, num_layers))

sources

Sources: recipes that resolve to concrete steering elements for a given model.

This module provides the ArtifactSource protocol with its fit recipes (ContrastiveFit, SinglePairFit) and the gate source ConditionPointSearch. A transform holds either a concrete artifact (a SteeringVector or a per-layer directions mapping) or a source; an Intervention's gate slot holds either a concrete gate or a gate source. Intervention.bind resolves sources, so steer-time computations (fits, searches) have a declarative home. resolve returns a defensive clone, and the underlying fit is memoized per model.

ArtifactSource

Bases: Protocol

A recipe for obtaining a steering artifact for a specific model.

Implementations MUST return a defensive clone from resolve (callers may move/mutate their copy) and SHOULD memoize the underlying fit per model so repeated resolves against one model (e.g., a parameter sweep) fit only once. Implementations whose fitted directions are positional ([T, H] with T > 1) declare it with a class-level produces_positional = True, which transform factories read for kind planning before the fit runs.

Source code in steerability/algorithms/state_control/common/sources.py

@runtime_checkable
class ArtifactSource(Protocol):
    """A recipe for obtaining a steering artifact for a specific model.

    Implementations MUST return a defensive clone from `resolve` (callers may move/mutate their
    copy) and SHOULD memoize the underlying fit per model so repeated resolves against one model
    (e.g., a parameter sweep) fit only once. Implementations whose fitted directions are
    positional (`[T, H]` with `T > 1`) declare it with a class-level `produces_positional =
    True`, which transform factories read for kind planning before the fit runs.
    """

    def resolve(
        self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, *, session: "SteeringSession | None" = None
    ) -> SteeringVector:
        """Return the steering artifact for this model (a fresh clone each call).

        `session` is a `SteeringSession` a capture-backed fit may extract hidden states
        through; sources whose fit requires a live model ignore it.
        """
        ...

resolve(model, tokenizer, *, session=None)

Return the steering artifact for this model (a fresh clone each call).

session is a SteeringSession a capture-backed fit may extract hidden states through; sources whose fit requires a live model ignore it.

Source code in steerability/algorithms/state_control/common/sources.py

def resolve(
    self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, *, session: "SteeringSession | None" = None
) -> SteeringVector:
    """Return the steering artifact for this model (a fresh clone each call).

    `session` is a `SteeringSession` a capture-backed fit may extract hidden states
    through; sources whose fit requires a live model ignore it.
    """
    ...

ConditionPointSearch dataclass

A gate recipe: contrastive condition data plus how to find the gate point.

Occupies an Intervention's gate slot. resolve_gate fits (or clones) the condition vector, resolves the condition point by grid search when search enables it and no manual layers are given, and assembles the gate: projected-cosine evidence over the condition directions at the resolved layers, decided by a per-layer threshold rule. An unconditional configuration (no condition vector or no resolved point) resolves to None.

The projected-cosine readout and the per-layer threshold rule both have wire forms, so an intervention gated this way lowers to intervention-capable backends (wire_readouts/wire_rules declare the kinds for pre-steer support checks).

Attributes:

Name	Type	Description
condition_vector	SteeringVector | None	Precomputed condition directions, cloned rather than refit.
condition_data	ContrastivePairs | dict | None	Contrastive pairs used to fit the condition vector and calibrate the search.
condition_fit	VectorTrainSpec	Fit configuration for the condition vector.
search	ConditionSearchSpec	Condition point search configuration.
layer_ids	Sequence[int] | None	Manual condition layers; disables the search when set.
threshold	float | None	Manual gate threshold, required with layer_ids.
comparator	Comparator	Gate comparator for the manual point ("ge" opens when score >= threshold, "le" when score <= threshold).
comparison_mode	CompMode	Runtime token aggregation for condition scoring.
resolved_point	dict | None	The (layer_ids, threshold, comparator) the last resolve produced, or None before resolution or for unconditional configurations.

Source code in steerability/algorithms/state_control/common/sources.py

@dataclass
class ConditionPointSearch:
    """A gate recipe: contrastive condition data plus how to find the gate point.

    Occupies an `Intervention`'s gate slot. `resolve_gate` fits (or clones) the condition
    vector, resolves the condition point by grid search when `search` enables it and no manual
    layers are given, and assembles the gate: projected-cosine evidence over the condition
    directions at the resolved layers, decided by a per-layer threshold rule. An unconditional
    configuration (no condition vector or no resolved point) resolves to None.

    The projected-cosine readout and the per-layer threshold rule both have wire forms, so an
    intervention gated this way lowers to intervention-capable backends
    (`wire_readouts`/`wire_rules` declare the kinds for pre-steer support checks).

    Attributes:
        condition_vector: Precomputed condition directions, cloned rather than refit.
        condition_data: Contrastive pairs used to fit the condition vector and calibrate the
            search.
        condition_fit: Fit configuration for the condition vector.
        search: Condition point search configuration.
        layer_ids: Manual condition layers; disables the search when set.
        threshold: Manual gate threshold, required with `layer_ids`.
        comparator: Gate comparator for the manual point (`"ge"` opens when score >= threshold,
            `"le"` when score <= threshold).
        comparison_mode: Runtime token aggregation for condition scoring.
        resolved_point: The `(layer_ids, threshold, comparator)` the last resolve produced, or
            None before resolution or for unconditional configurations.
    """

    wire_readouts: ClassVar[frozenset[str] | None] = frozenset({"projected_cosine"})
    wire_rules: ClassVar[frozenset[str] | None] = frozenset({"per_key_threshold"})
    access: ClassVar[ModelAccess] = ModelAccess.MODULE
    artifact_class: ClassVar[str] = "calibrated"

    condition_vector: SteeringVector | None = None
    condition_data: ContrastivePairs | dict | None = None
    condition_fit: VectorTrainSpec = field(
        default_factory=lambda: VectorTrainSpec(prompt_format="chat_prompt", location="layer_input")
    )
    search: ConditionSearchSpec = field(default_factory=ConditionSearchSpec)
    layer_ids: Sequence[int] | None = None
    threshold: float | None = None
    comparator: Comparator = "ge"
    comparison_mode: CompMode = "mean"

    resolved_point: dict | None = field(default=None, init=False, repr=False, compare=False)

    def __post_init__(self):
        if self.condition_data is not None and not isinstance(self.condition_data, ContrastivePairs):
            self.condition_data = as_contrastive_pairs(self.condition_data)
        if self.comparator not in ("ge", "le"):
            raise ValueError(f"comparator must be 'ge' or 'le'; got {self.comparator!r}.")

    def resolve_gate(
        self,
        model: PreTrainedModel | None,
        tokenizer: PreTrainedTokenizerBase,
        *,
        layout: "ModelFacts | None" = None,
        session: "SteeringSession | None" = None,
    ) -> "Gate | None":
        """Resolve the gate for `model`.

        Args:
            model: The model to search against; required when the search runs or the condition
                vector is fitted from data.
            tokenizer: Tokenizer used to encode the condition data.
            layout: Structural facts, used to bounds-check condition layers.
            session: Optional `SteeringSession` for capture-backed fitting and calibration.

        Returns:
            The gate, or None for unconditional configurations.

        Raises:
            ValueError: If a manual threshold is set without a condition vector, or a condition
                layer lacks a direction.
        """
        from steerability.algorithms.state_control.common.gating import (
            Evidence,
            Gate,
            PerKeyThreshold,
            ProjectedCosineReadout,
        )
        from steerability.algorithms.state_control.common.selectors import ConditionPointSelector

        condition_vec = self.condition_vector.clone() if self.condition_vector is not None else None
        condition_supplied = condition_vec is not None or self.condition_data is not None
        if condition_supplied and condition_vec is None:
            if self.condition_fit.method == "mean_diff" and self.condition_fit.accumulate == "suffix-only":
                raise ValueError(
                    "method='mean_diff' does not support accumulate='suffix-only'; "
                    "use accumulate='all' or 'last_token', or method='pca_pairwise'/'pca_center'."
                )
            estimator = (
                MeanDifferenceEstimator()
                if self.condition_fit.method == "mean_diff"
                else ContrastiveDirectionEstimator()
            )
            condition_vec = estimator.fit(
                model, tokenizer, data=self.condition_data, spec=self.condition_fit, session=session,
            )
            if model is not None:
                device = next(model.parameters()).device
                condition_vec = condition_vec.to(device, dtype=model.dtype)

        layer_ids = self.layer_ids
        threshold = self.threshold
        comparator = self.comparator

        if condition_supplied and condition_vec is not None:
            if self.search.auto_find and layer_ids is None and self.condition_data is not None:
                result = ConditionPointSelector().select(
                    model=model,
                    tokenizer=tokenizer,
                    condition_directions=condition_vec.directions,
                    data=self.condition_data,
                    fit_spec=self.condition_fit,
                    search_spec=self.search,
                    comparison_mode=self.comparison_mode,
                    session=session,
                )
                layer_ids = [result.layer_id]
                threshold = result.threshold
                comparator = result.comparator

        layer_set = sorted(set(int(lid) for lid in (layer_ids or [])))
        conditional = bool(layer_set) and threshold is not None
        if conditional and condition_vec is None:
            raise ValueError("Conditional gating requires a condition vector.")
        if not conditional:
            self.resolved_point = None
            return None

        missing = [lid for lid in layer_set if lid not in condition_vec.directions]
        if missing:
            raise ValueError(f"Condition vector has no direction for condition layer(s) {missing}.")

        readout = ProjectedCosineReadout(
            {lid: condition_vec.directions[lid] for lid in layer_set},
        )
        rule = PerKeyThreshold(threshold=threshold, comparator=comparator, aggregate="any")
        self.resolved_point = {
            "layer_ids": layer_set,
            "threshold": threshold,
            "comparator": comparator,
            "comparison_mode": self.comparison_mode,
        }
        return Gate(
            Evidence(tuple(layer_set), readout, pooling=self.comparison_mode),
            rule,
        )

access = ModelAccess.MODULE class-attribute

artifact_class = 'calibrated' class-attribute

comparator = 'ge' class-attribute instance-attribute

comparison_mode = 'mean' class-attribute instance-attribute

condition_data = None class-attribute instance-attribute

condition_fit = field(default_factory=lambda: VectorTrainSpec(prompt_format='chat_prompt', location='layer_input')) class-attribute instance-attribute

condition_vector = None class-attribute instance-attribute

layer_ids = None class-attribute instance-attribute

resolved_point = field(default=None, init=False, repr=False, compare=False) class-attribute instance-attribute

search = field(default_factory=ConditionSearchSpec) class-attribute instance-attribute

threshold = None class-attribute instance-attribute

wire_readouts = frozenset({'projected_cosine'}) class-attribute

wire_rules = frozenset({'per_key_threshold'}) class-attribute

__init__(condition_vector=None, condition_data=None, condition_fit=(lambda: VectorTrainSpec(prompt_format='chat_prompt', location='layer_input'))(), search=ConditionSearchSpec(), layer_ids=None, threshold=None, comparator='ge', comparison_mode='mean')

__post_init__()

Source code in steerability/algorithms/state_control/common/sources.py

def __post_init__(self):
    if self.condition_data is not None and not isinstance(self.condition_data, ContrastivePairs):
        self.condition_data = as_contrastive_pairs(self.condition_data)
    if self.comparator not in ("ge", "le"):
        raise ValueError(f"comparator must be 'ge' or 'le'; got {self.comparator!r}.")

resolve_gate(model, tokenizer, *, layout=None, session=None)

Resolve the gate for model.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	The model to search against; required when the search runs or the condition vector is fitted from data.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer used to encode the condition data.	required
layout	'ModelFacts | None'	Structural facts, used to bounds-check condition layers.	None
session	'SteeringSession | None'	Optional SteeringSession for capture-backed fitting and calibration.	None

Returns:

Type	Description
'Gate | None'	The gate, or None for unconditional configurations.

Raises:

Type	Description
ValueError	If a manual threshold is set without a condition vector, or a condition layer lacks a direction.

Source code in steerability/algorithms/state_control/common/sources.py

def resolve_gate(
    self,
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase,
    *,
    layout: "ModelFacts | None" = None,
    session: "SteeringSession | None" = None,
) -> "Gate | None":
    """Resolve the gate for `model`.

    Args:
        model: The model to search against; required when the search runs or the condition
            vector is fitted from data.
        tokenizer: Tokenizer used to encode the condition data.
        layout: Structural facts, used to bounds-check condition layers.
        session: Optional `SteeringSession` for capture-backed fitting and calibration.

    Returns:
        The gate, or None for unconditional configurations.

    Raises:
        ValueError: If a manual threshold is set without a condition vector, or a condition
            layer lacks a direction.
    """
    from steerability.algorithms.state_control.common.gating import (
        Evidence,
        Gate,
        PerKeyThreshold,
        ProjectedCosineReadout,
    )
    from steerability.algorithms.state_control.common.selectors import ConditionPointSelector

    condition_vec = self.condition_vector.clone() if self.condition_vector is not None else None
    condition_supplied = condition_vec is not None or self.condition_data is not None
    if condition_supplied and condition_vec is None:
        if self.condition_fit.method == "mean_diff" and self.condition_fit.accumulate == "suffix-only":
            raise ValueError(
                "method='mean_diff' does not support accumulate='suffix-only'; "
                "use accumulate='all' or 'last_token', or method='pca_pairwise'/'pca_center'."
            )
        estimator = (
            MeanDifferenceEstimator()
            if self.condition_fit.method == "mean_diff"
            else ContrastiveDirectionEstimator()
        )
        condition_vec = estimator.fit(
            model, tokenizer, data=self.condition_data, spec=self.condition_fit, session=session,
        )
        if model is not None:
            device = next(model.parameters()).device
            condition_vec = condition_vec.to(device, dtype=model.dtype)

    layer_ids = self.layer_ids
    threshold = self.threshold
    comparator = self.comparator

    if condition_supplied and condition_vec is not None:
        if self.search.auto_find and layer_ids is None and self.condition_data is not None:
            result = ConditionPointSelector().select(
                model=model,
                tokenizer=tokenizer,
                condition_directions=condition_vec.directions,
                data=self.condition_data,
                fit_spec=self.condition_fit,
                search_spec=self.search,
                comparison_mode=self.comparison_mode,
                session=session,
            )
            layer_ids = [result.layer_id]
            threshold = result.threshold
            comparator = result.comparator

    layer_set = sorted(set(int(lid) for lid in (layer_ids or [])))
    conditional = bool(layer_set) and threshold is not None
    if conditional and condition_vec is None:
        raise ValueError("Conditional gating requires a condition vector.")
    if not conditional:
        self.resolved_point = None
        return None

    missing = [lid for lid in layer_set if lid not in condition_vec.directions]
    if missing:
        raise ValueError(f"Condition vector has no direction for condition layer(s) {missing}.")

    readout = ProjectedCosineReadout(
        {lid: condition_vec.directions[lid] for lid in layer_set},
    )
    rule = PerKeyThreshold(threshold=threshold, comparator=comparator, aggregate="any")
    self.resolved_point = {
        "layer_ids": layer_set,
        "threshold": threshold,
        "comparator": comparator,
        "comparison_mode": self.comparison_mode,
    }
    return Gate(
        Evidence(tuple(layer_set), readout, pooling=self.comparison_mode),
        rule,
    )

ContrastiveFit dataclass

A fit recipe: contrastive pairs plus how to extract a per-layer direction.

The five spec fields (method, accumulate, batch_size, prompt_format, location) drive the built-in estimators. "mean_diff" dispatches to MeanDifferenceEstimator, everything else to ContrastiveDirectionEstimator, and the fitted artifact records its extraction boundary as meta["location"], which Intervention.bind checks against the consuming intervention's boundary. When a custom estimator is supplied, the spec fields are ignored (a warning is emitted), no location is recorded, and fitting delegates to estimator.fit(model, tokenizer, data=<coerced pairs>, **(estimator_kwargs or {})).

The fitted master vector is memoized in a single weakref slot keyed by model identity. The same model resolved repeatedly fits once, a different model refits, and alternating models (A→B→A) refit on each switch. Every resolve returns an independent clone, so a consumer may freely .to(...) or mutate its copy without touching the master.

Attributes:

Name	Type	Description
data	ContrastivePairs | dict	Contrastive pairs (or a dict coerced via as_contrastive_pairs).
method	str	Direction-extraction method used to fit the per-layer direction.
accumulate	str	Hidden-state span selection for the fit.
batch_size	int	Forward-pass batch size used during fitting.
prompt_format	PromptFormat	How to render pairs into model-ready text.
location	HiddenStateLocation	Residual-stream boundary to fit at.
normalize	bool	L2-normalize the fitted master per layer once, before caching.
estimator	BaseEstimator | None	Optional custom BaseEstimator; when set, the spec fields are ignored.
estimator_kwargs	dict | None	Extra kwargs forwarded to a custom estimator.fit(...).

Source code in steerability/algorithms/state_control/common/sources.py

@dataclass
class ContrastiveFit:
    """A fit recipe: contrastive pairs plus how to extract a per-layer direction.

    The five spec fields (`method`, `accumulate`, `batch_size`, `prompt_format`, `location`) drive
    the built-in estimators. `"mean_diff"` dispatches to `MeanDifferenceEstimator`, everything else
    to `ContrastiveDirectionEstimator`, and the fitted artifact records its extraction boundary as
    `meta["location"]`, which `Intervention.bind` checks against the consuming intervention's
    boundary. When a custom `estimator` is supplied, the spec fields are ignored (a warning is
    emitted), no location is recorded, and fitting delegates to
    `estimator.fit(model, tokenizer, data=<coerced pairs>, **(estimator_kwargs or {}))`.

    The fitted master vector is memoized in a single weakref slot keyed by model identity. The same
    model resolved repeatedly fits once, a different model refits, and alternating models (A→B→A)
    refit on each switch. Every `resolve` returns an independent clone, so a consumer may freely
    `.to(...)` or mutate its copy without touching the master.

    Attributes:
        data: Contrastive pairs (or a dict coerced via `as_contrastive_pairs`).
        method: Direction-extraction method used to fit the per-layer direction.
        accumulate: Hidden-state span selection for the fit.
        batch_size: Forward-pass batch size used during fitting.
        prompt_format: How to render pairs into model-ready text.
        location: Residual-stream boundary to fit at.
        normalize: L2-normalize the fitted master per layer once, before caching.
        estimator: Optional custom `BaseEstimator`; when set, the spec fields are ignored.
        estimator_kwargs: Extra kwargs forwarded to a custom `estimator.fit(...)`.
    """

    produces_positional: ClassVar[bool] = False
    artifact_class: ClassVar[str] = "direction"

    data: ContrastivePairs | dict
    method: str = "pca_pairwise"
    accumulate: str = "all"
    batch_size: int = 8
    prompt_format: PromptFormat = "chat_completion"
    location: HiddenStateLocation = "layer_output"
    normalize: bool = False
    estimator: BaseEstimator | None = None
    estimator_kwargs: dict | None = None

    _model_ref: "weakref.ref | None" = field(default=None, init=False, repr=False, compare=False)
    _master: SteeringVector | None = field(default=None, init=False, repr=False, compare=False)

    @property
    def access(self) -> ModelAccess:
        """`ModelAccess.CAPTURE` for the built-in estimators, whose extraction runs through
        session capture; a custom estimator may need the live model, so it declares
        `ModelAccess.MODULE`."""
        return ModelAccess.MODULE if self.estimator is not None else ModelAccess.CAPTURE

    def __post_init__(self):
        if not isinstance(self.data, ContrastivePairs):
            self.data = as_contrastive_pairs(self.data)

        spec_customized = (
            self.method != "pca_pairwise"
            or self.accumulate != "all"
            or self.batch_size != 8
            or self.prompt_format != "chat_completion"
            or self.location != "layer_output"
        )
        if self.estimator is not None and spec_customized:
            warnings.warn(
                "method/accumulate/batch_size/prompt_format/location are ignored when a custom "
                "estimator is supplied; the estimator owns its config via estimator_kwargs.",
                UserWarning,
            )
        if self.estimator is None and self.estimator_kwargs is not None:
            warnings.warn("estimator_kwargs is inert without a custom estimator.", UserWarning)

    def _fit(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, session=None) -> SteeringVector:
        """Fit the master steering vector (no caching, no cloning)."""
        if self.estimator is not None:
            try:
                master = self.estimator.fit(
                    model, tokenizer, data=self.data, session=session, **(self.estimator_kwargs or {})
                )
            except TypeError:
                master = self.estimator.fit(model, tokenizer, data=self.data, **(self.estimator_kwargs or {}))
        else:
            spec = VectorTrainSpec(
                method=self.method,
                accumulate=self.accumulate,
                batch_size=self.batch_size,
                prompt_format=self.prompt_format,
                location=self.location,
            )
            estimator = MeanDifferenceEstimator() if self.method == "mean_diff" else ContrastiveDirectionEstimator()
            master = estimator.fit(model, tokenizer, data=self.data, spec=spec, session=session)
            master.meta["location"] = self.location

        if self.normalize:
            master = master.normalized()
        return master

    def resolve(
        self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, *, session: "SteeringSession | None" = None
    ) -> SteeringVector:
        """Return a fresh clone of the fitted artifact for `model`, fitting once and memoizing.

        Args:
            model: The model to fit against (or a memo hit for the same model), or None to
                fit through `session` capture.
            tokenizer: Tokenizer used to encode the contrastive pairs when fitting.
            session: Optional `SteeringSession` the estimator extracts hidden states through
                when `model` is None (capture-backed fitting).

        Returns:
            An independent `SteeringVector` clone the caller owns.
        """
        if model is not None and self._model_ref is not None and self._model_ref() is model \
                and self._master is not None:
            return self._master.clone()
        master = self._fit(model, tokenizer, session=session)
        if model is not None:
            self._model_ref = weakref.ref(model)
            self._master = master
        return master.clone()

_master = field(default=None, init=False, repr=False, compare=False) class-attribute instance-attribute

_model_ref = field(default=None, init=False, repr=False, compare=False) class-attribute instance-attribute

access property

ModelAccess.CAPTURE for the built-in estimators, whose extraction runs through session capture; a custom estimator may need the live model, so it declares ModelAccess.MODULE.

accumulate = 'all' class-attribute instance-attribute

artifact_class = 'direction' class-attribute

batch_size = 8 class-attribute instance-attribute

data instance-attribute

estimator = None class-attribute instance-attribute

estimator_kwargs = None class-attribute instance-attribute

location = 'layer_output' class-attribute instance-attribute

method = 'pca_pairwise' class-attribute instance-attribute

normalize = False class-attribute instance-attribute

produces_positional = False class-attribute

prompt_format = 'chat_completion' class-attribute instance-attribute

__init__(data, method='pca_pairwise', accumulate='all', batch_size=8, prompt_format='chat_completion', location='layer_output', normalize=False, estimator=None, estimator_kwargs=None)

__post_init__()

Source code in steerability/algorithms/state_control/common/sources.py

def __post_init__(self):
    if not isinstance(self.data, ContrastivePairs):
        self.data = as_contrastive_pairs(self.data)

    spec_customized = (
        self.method != "pca_pairwise"
        or self.accumulate != "all"
        or self.batch_size != 8
        or self.prompt_format != "chat_completion"
        or self.location != "layer_output"
    )
    if self.estimator is not None and spec_customized:
        warnings.warn(
            "method/accumulate/batch_size/prompt_format/location are ignored when a custom "
            "estimator is supplied; the estimator owns its config via estimator_kwargs.",
            UserWarning,
        )
    if self.estimator is None and self.estimator_kwargs is not None:
        warnings.warn("estimator_kwargs is inert without a custom estimator.", UserWarning)

_fit(model, tokenizer, session=None)

Fit the master steering vector (no caching, no cloning).

Source code in steerability/algorithms/state_control/common/sources.py

def _fit(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, session=None) -> SteeringVector:
    """Fit the master steering vector (no caching, no cloning)."""
    if self.estimator is not None:
        try:
            master = self.estimator.fit(
                model, tokenizer, data=self.data, session=session, **(self.estimator_kwargs or {})
            )
        except TypeError:
            master = self.estimator.fit(model, tokenizer, data=self.data, **(self.estimator_kwargs or {}))
    else:
        spec = VectorTrainSpec(
            method=self.method,
            accumulate=self.accumulate,
            batch_size=self.batch_size,
            prompt_format=self.prompt_format,
            location=self.location,
        )
        estimator = MeanDifferenceEstimator() if self.method == "mean_diff" else ContrastiveDirectionEstimator()
        master = estimator.fit(model, tokenizer, data=self.data, spec=spec, session=session)
        master.meta["location"] = self.location

    if self.normalize:
        master = master.normalized()
    return master

resolve(model, tokenizer, *, session=None)

Return a fresh clone of the fitted artifact for model, fitting once and memoizing.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	The model to fit against (or a memo hit for the same model), or None to fit through session capture.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer used to encode the contrastive pairs when fitting.	required
session	'SteeringSession | None'	Optional SteeringSession the estimator extracts hidden states through when model is None (capture-backed fitting).	None

Returns:

Type	Description
SteeringVector	An independent SteeringVector clone the caller owns.

Source code in steerability/algorithms/state_control/common/sources.py

def resolve(
    self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, *, session: "SteeringSession | None" = None
) -> SteeringVector:
    """Return a fresh clone of the fitted artifact for `model`, fitting once and memoizing.

    Args:
        model: The model to fit against (or a memo hit for the same model), or None to
            fit through `session` capture.
        tokenizer: Tokenizer used to encode the contrastive pairs when fitting.
        session: Optional `SteeringSession` the estimator extracts hidden states through
            when `model` is None (capture-backed fitting).

    Returns:
        An independent `SteeringVector` clone the caller owns.
    """
    if model is not None and self._model_ref is not None and self._model_ref() is model \
            and self._master is not None:
        return self._master.clone()
    master = self._fit(model, tokenizer, session=session)
    if model is not None:
        self._model_ref = weakref.ref(model)
        self._master = master
    return master.clone()

LayerFilteredFit dataclass

Wraps a source and restricts the resolved directions to a layer range.

Access, artifact class, and positional-ness delegate to the wrapped source. The filtered result keeps the inner artifact's metadata and per-layer statistics for the surviving layers.

Attributes:

Name	Type	Description
inner	'ArtifactSource'	The wrapped source.
layer_range	tuple[int, int] | None	0-based half-open (start, end) range; None passes every layer through.

Source code in steerability/algorithms/state_control/common/sources.py

@dataclass
class LayerFilteredFit:
    """Wraps a source and restricts the resolved directions to a layer range.

    Access, artifact class, and positional-ness delegate to the wrapped source. The filtered
    result keeps the inner artifact's metadata and per-layer statistics for the surviving
    layers.

    Attributes:
        inner: The wrapped source.
        layer_range: 0-based half-open `(start, end)` range; None passes every layer through.
    """

    inner: "ArtifactSource"
    layer_range: tuple[int, int] | None = None

    @property
    def access(self) -> ModelAccess | None:
        return getattr(self.inner, "access", None)

    @property
    def artifact_class(self) -> str | None:
        return getattr(self.inner, "artifact_class", None)

    @property
    def produces_positional(self) -> bool:
        return bool(getattr(self.inner, "produces_positional", False))

    def resolve(
        self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, *, session: "SteeringSession | None" = None
    ) -> SteeringVector:
        """Resolve the inner source and filter its directions to `layer_range`."""
        resolved = self.inner.resolve(model, tokenizer, session=session)
        if self.layer_range is None:
            return resolved
        start, end = self.layer_range
        directions = {
            layer_id: direction
            for layer_id, direction in resolved.directions.items()
            if start <= layer_id < end
        }
        return SteeringVector(
            model_type=resolved.model_type,
            directions=directions,
            num_heads=resolved.num_heads,
            head_dim=resolved.head_dim,
            explained_variances=resolved.explained_variances,
            probe_accuracies=resolved.probe_accuracies,
            meta=dict(resolved.meta),
        )

access property

artifact_class property

inner instance-attribute

layer_range = None class-attribute instance-attribute

produces_positional property

__init__(inner, layer_range=None)

resolve(model, tokenizer, *, session=None)

Resolve the inner source and filter its directions to layer_range.

Source code in steerability/algorithms/state_control/common/sources.py

def resolve(
    self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, *, session: "SteeringSession | None" = None
) -> SteeringVector:
    """Resolve the inner source and filter its directions to `layer_range`."""
    resolved = self.inner.resolve(model, tokenizer, session=session)
    if self.layer_range is None:
        return resolved
    start, end = self.layer_range
    directions = {
        layer_id: direction
        for layer_id, direction in resolved.directions.items()
        if start <= layer_id < end
    }
    return SteeringVector(
        model_type=resolved.model_type,
        directions=directions,
        num_heads=resolved.num_heads,
        head_dim=resolved.head_dim,
        explained_variances=resolved.explained_variances,
        probe_accuracies=resolved.probe_accuracies,
        meta=dict(resolved.meta),
    )

SinglePairFit dataclass

A fit recipe for a positional steering vector: per-token differences of one contrast pair.

Produces [T, H] directions per layer (the ActAdd extraction), captured at the layer-input boundary and recorded as meta["location"] = "layer_input". Fitting requires a live model, since the pair is tokenized and co-padded as one batch and remote capture serves per-prompt layouts that do not preserve the co-padded positional structure. The fitted master is memoized per model; every resolve returns an independent clone.

Attributes:

Name	Type	Description
positive_prompt	str	The steering-direction prompt.
negative_prompt	str	The contrast prompt.
normalize	bool	L2-normalize each per-position vector once, before caching.

Source code in steerability/algorithms/state_control/common/sources.py

@dataclass
class SinglePairFit:
    """A fit recipe for a positional steering vector: per-token differences of one contrast pair.

    Produces `[T, H]` directions per layer (the ActAdd extraction), captured at the layer-input
    boundary and recorded as `meta["location"] = "layer_input"`. Fitting requires a live model,
    since the pair is tokenized and co-padded as one batch and remote capture serves per-prompt
    layouts that do not preserve the co-padded positional structure. The fitted master is
    memoized per model; every `resolve` returns an independent clone.

    Attributes:
        positive_prompt: The steering-direction prompt.
        negative_prompt: The contrast prompt.
        normalize: L2-normalize each per-position vector once, before caching.
    """

    produces_positional: ClassVar[bool] = True
    access: ClassVar[ModelAccess] = ModelAccess.MODULE
    artifact_class: ClassVar[str] = "direction"

    positive_prompt: str
    negative_prompt: str
    normalize: bool = False

    _model_ref: "weakref.ref | None" = field(default=None, init=False, repr=False, compare=False)
    _master: SteeringVector | None = field(default=None, init=False, repr=False, compare=False)

    def resolve(
        self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, *, session: "SteeringSession | None" = None
    ) -> SteeringVector:
        """Return a fresh clone of the fitted artifact for `model`, fitting once and memoizing.

        Args:
            model: The model to fit against (or a memo hit for the same model).
            tokenizer: Tokenizer used to encode the prompt pair.
            session: Ignored; the positional fit requires a live model.

        Returns:
            An independent `SteeringVector` clone the caller owns.

        Raises:
            ValueError: If `model` is None.
        """
        if model is None:
            raise ValueError("Fitting ActAdd from a prompt pair requires a live model at steer time.")
        if self._model_ref is not None and self._model_ref() is model and self._master is not None:
            return self._master.clone()

        from steerability.algorithms.state_control.common.estimators import SinglePairEstimator

        master = SinglePairEstimator().fit(
            model, tokenizer,
            positive_prompt=self.positive_prompt,
            negative_prompt=self.negative_prompt,
        )
        if self.normalize:
            master = master.clone()
            for layer_id, direction in master.directions.items():
                norms = direction.norm(dim=-1, keepdim=True)
                master.directions[layer_id] = direction / (norms + 1e-8)
        self._model_ref = weakref.ref(model)
        self._master = master
        return master.clone()

_master = field(default=None, init=False, repr=False, compare=False) class-attribute instance-attribute

_model_ref = field(default=None, init=False, repr=False, compare=False) class-attribute instance-attribute

access = ModelAccess.MODULE class-attribute

artifact_class = 'direction' class-attribute

negative_prompt instance-attribute

normalize = False class-attribute instance-attribute

positive_prompt instance-attribute

produces_positional = True class-attribute

__init__(positive_prompt, negative_prompt, normalize=False)

resolve(model, tokenizer, *, session=None)

Return a fresh clone of the fitted artifact for model, fitting once and memoizing.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	The model to fit against (or a memo hit for the same model).	required
tokenizer	PreTrainedTokenizerBase	Tokenizer used to encode the prompt pair.	required
session	'SteeringSession | None'	Ignored; the positional fit requires a live model.	None

Returns:

Type	Description
SteeringVector	An independent SteeringVector clone the caller owns.

Raises:

Type	Description
ValueError	If model is None.

Source code in steerability/algorithms/state_control/common/sources.py

def resolve(
    self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, *, session: "SteeringSession | None" = None
) -> SteeringVector:
    """Return a fresh clone of the fitted artifact for `model`, fitting once and memoizing.

    Args:
        model: The model to fit against (or a memo hit for the same model).
        tokenizer: Tokenizer used to encode the prompt pair.
        session: Ignored; the positional fit requires a live model.

    Returns:
        An independent `SteeringVector` clone the caller owns.

    Raises:
        ValueError: If `model` is None.
    """
    if model is None:
        raise ValueError("Fitting ActAdd from a prompt pair requires a live model at steer time.")
    if self._model_ref is not None and self._model_ref() is model and self._master is not None:
        return self._master.clone()

    from steerability.algorithms.state_control.common.estimators import SinglePairEstimator

    master = SinglePairEstimator().fit(
        model, tokenizer,
        positive_prompt=self.positive_prompt,
        negative_prompt=self.negative_prompt,
    )
    if self.normalize:
        master = master.clone()
        for layer_id, direction in master.directions.items():
            norms = direction.norm(dim=-1, keepdim=True)
            master.directions[layer_id] = direction / (norms + 1e-8)
    self._model_ref = weakref.ref(model)
    self._master = master
    return master.clone()

VerifiedPrecomputed

A precomputed source that enforces recorded provenance at resolve time.

Wraps a concrete SteeringVector together with the fingerprints of the side that produced it. Resolution is model-free (ModelAccess.FACTS) and reads structural facts from the session layout (or the live model when one is given): the vector's model_type and width are always checked, and the recorded model fingerprint is checked per the policy. A "calibrated" artifact on a mismatched model raises under policy="strict" and warns under "warn"; a "direction" artifact warns under both, since transferring a direction across fine-tunes of one architecture is a deliberate act. policy="off" skips every check. Checks that cannot run (no layout available, or no recorded value) are skipped.

Parameters:

Name	Type	Description	Default
steering_vector	SteeringVector	The concrete artifact.	required
provenance	Mapping | None	Producing-side fingerprints (mapping or ArtifactProvenance-shaped), with keys model_fingerprint, backend_spec_hash, tokenizer_fingerprint.	None
artifact_class	str	"direction" or "calibrated".	'direction'
policy	str	"strict", "warn", or "off".	'strict'

Source code in steerability/algorithms/state_control/common/sources.py

class VerifiedPrecomputed:
    """A precomputed source that enforces recorded provenance at resolve time.

    Wraps a concrete `SteeringVector` together with the fingerprints of the side that produced
    it. Resolution is model-free (`ModelAccess.FACTS`) and reads structural facts from the
    session layout (or the live model when one is given): the vector's `model_type` and width
    are always checked, and the recorded model fingerprint is checked per the policy. A
    `"calibrated"` artifact on a mismatched model raises under `policy="strict"` and warns
    under `"warn"`; a `"direction"` artifact warns under both, since transferring a direction
    across fine-tunes of one architecture is a deliberate act. `policy="off"` skips every
    check. Checks that cannot run (no layout available, or no recorded value) are skipped.

    Args:
        steering_vector: The concrete artifact.
        provenance: Producing-side fingerprints (mapping or `ArtifactProvenance`-shaped), with
            keys `model_fingerprint`, `backend_spec_hash`, `tokenizer_fingerprint`.
        artifact_class: `"direction"` or `"calibrated"`.
        policy: `"strict"`, `"warn"`, or `"off"`.
    """

    access: ClassVar[ModelAccess] = ModelAccess.FACTS

    def __init__(
        self,
        steering_vector: SteeringVector,
        provenance: Mapping | None = None,
        artifact_class: str = "direction",
        policy: str = "strict",
    ):
        if policy not in ("strict", "warn", "off"):
            raise ValueError(f"policy must be 'strict', 'warn', or 'off'; got {policy!r}.")
        self.steering_vector = steering_vector
        self.provenance = dict(provenance or {})
        self.artifact_class = artifact_class
        self.policy = policy

    @property
    def produces_positional(self) -> bool:
        return self.steering_vector.is_positional

    def _report(self, message: str, *, hard: bool) -> None:
        if self.policy == "off":
            return
        if hard and self.policy == "strict":
            raise ValueError(message)
        warnings.warn(message, UserWarning)

    def _check(self, layout) -> None:
        vector = self.steering_vector
        if layout is None:
            return

        layout_type = getattr(layout, "model_type", None)
        if vector.model_type not in ("unknown", None) and layout_type not in (None, "unknown") \
                and vector.model_type != layout_type:
            self._report(
                f"Precomputed steering artifact was produced for model_type "
                f"{vector.model_type!r} but this pipeline serves {layout_type!r}.",
                hard=True,
            )

        if vector.num_heads is not None and vector.head_dim is not None:
            num_heads = getattr(layout, "num_attention_heads", None)
            head_dim = getattr(layout, "head_dim", None)
            if (num_heads is not None and vector.num_heads != num_heads) or \
                    (head_dim is not None and vector.head_dim != head_dim):
                self._report(
                    f"Precomputed per-head artifact has num_heads={vector.num_heads}, "
                    f"head_dim={vector.head_dim} but the model has num_heads={num_heads}, "
                    f"head_dim={head_dim}.",
                    hard=True,
                )
        elif vector.directions:
            width = next(iter(vector.directions.values())).size(-1)
            hidden_size = getattr(layout, "hidden_size", None)
            if hidden_size and width != hidden_size:
                self._report(
                    f"Precomputed steering artifact has width {width} but the model's hidden "
                    f"size is {hidden_size}.",
                    hard=True,
                )

        recorded = self.provenance.get("model_fingerprint")
        live = getattr(layout, "model_fingerprint", None)
        if recorded and live and recorded != live:
            self._report(
                f"Precomputed {self.artifact_class} artifact was produced on a different "
                f"model (fingerprint {recorded!r} vs {live!r})."
                + ("" if self.artifact_class == "calibrated"
                   else " Direction artifacts may transfer across fine-tunes of one "
                        "architecture; verify the behavior."),
                hard=(self.artifact_class == "calibrated"),
            )

    def resolve(
        self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, *, session: "SteeringSession | None" = None
    ) -> SteeringVector:
        """Check provenance against the resolution venue and return a fresh clone.

        Args:
            model: The live model, or None when a session layout supplies the facts.
            tokenizer: Unused.
            session: Optional `SteeringSession` whose `layout` supplies structural facts.

        Returns:
            An independent `SteeringVector` clone the caller owns.

        Raises:
            ValueError: Under `policy="strict"`, if the model type or width mismatches, or a
                `"calibrated"` artifact's recorded fingerprint differs from the live model's.
        """
        layout = None
        if session is not None:
            try:
                layout = session.layout
            except Exception:
                layout = None
        if layout is None and model is not None:
            from steerability.algorithms.state_control.common.layout_facts import resolve_layout

            layout = resolve_layout(model, None)
        self._check(layout)
        return self.steering_vector.clone()

access = ModelAccess.FACTS class-attribute

artifact_class = artifact_class instance-attribute

policy = policy instance-attribute

produces_positional property

provenance = dict(provenance or {}) instance-attribute

steering_vector = steering_vector instance-attribute

__init__(steering_vector, provenance=None, artifact_class='direction', policy='strict')

Source code in steerability/algorithms/state_control/common/sources.py

def __init__(
    self,
    steering_vector: SteeringVector,
    provenance: Mapping | None = None,
    artifact_class: str = "direction",
    policy: str = "strict",
):
    if policy not in ("strict", "warn", "off"):
        raise ValueError(f"policy must be 'strict', 'warn', or 'off'; got {policy!r}.")
    self.steering_vector = steering_vector
    self.provenance = dict(provenance or {})
    self.artifact_class = artifact_class
    self.policy = policy

_check(layout)

Source code in steerability/algorithms/state_control/common/sources.py

def _check(self, layout) -> None:
    vector = self.steering_vector
    if layout is None:
        return

    layout_type = getattr(layout, "model_type", None)
    if vector.model_type not in ("unknown", None) and layout_type not in (None, "unknown") \
            and vector.model_type != layout_type:
        self._report(
            f"Precomputed steering artifact was produced for model_type "
            f"{vector.model_type!r} but this pipeline serves {layout_type!r}.",
            hard=True,
        )

    if vector.num_heads is not None and vector.head_dim is not None:
        num_heads = getattr(layout, "num_attention_heads", None)
        head_dim = getattr(layout, "head_dim", None)
        if (num_heads is not None and vector.num_heads != num_heads) or \
                (head_dim is not None and vector.head_dim != head_dim):
            self._report(
                f"Precomputed per-head artifact has num_heads={vector.num_heads}, "
                f"head_dim={vector.head_dim} but the model has num_heads={num_heads}, "
                f"head_dim={head_dim}.",
                hard=True,
            )
    elif vector.directions:
        width = next(iter(vector.directions.values())).size(-1)
        hidden_size = getattr(layout, "hidden_size", None)
        if hidden_size and width != hidden_size:
            self._report(
                f"Precomputed steering artifact has width {width} but the model's hidden "
                f"size is {hidden_size}.",
                hard=True,
            )

    recorded = self.provenance.get("model_fingerprint")
    live = getattr(layout, "model_fingerprint", None)
    if recorded and live and recorded != live:
        self._report(
            f"Precomputed {self.artifact_class} artifact was produced on a different "
            f"model (fingerprint {recorded!r} vs {live!r})."
            + ("" if self.artifact_class == "calibrated"
               else " Direction artifacts may transfer across fine-tunes of one "
                    "architecture; verify the behavior."),
            hard=(self.artifact_class == "calibrated"),
        )

_report(message, *, hard)

Source code in steerability/algorithms/state_control/common/sources.py

def _report(self, message: str, *, hard: bool) -> None:
    if self.policy == "off":
        return
    if hard and self.policy == "strict":
        raise ValueError(message)
    warnings.warn(message, UserWarning)

resolve(model, tokenizer, *, session=None)

Check provenance against the resolution venue and return a fresh clone.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	The live model, or None when a session layout supplies the facts.	required
tokenizer	PreTrainedTokenizerBase	Unused.	required
session	'SteeringSession | None'	Optional SteeringSession whose layout supplies structural facts.	None

Returns:

Type	Description
SteeringVector	An independent SteeringVector clone the caller owns.

Raises:

Type	Description
ValueError	Under policy="strict", if the model type or width mismatches, or a "calibrated" artifact's recorded fingerprint differs from the live model's.

Source code in steerability/algorithms/state_control/common/sources.py

def resolve(
    self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, *, session: "SteeringSession | None" = None
) -> SteeringVector:
    """Check provenance against the resolution venue and return a fresh clone.

    Args:
        model: The live model, or None when a session layout supplies the facts.
        tokenizer: Unused.
        session: Optional `SteeringSession` whose `layout` supplies structural facts.

    Returns:
        An independent `SteeringVector` clone the caller owns.

    Raises:
        ValueError: Under `policy="strict"`, if the model type or width mismatches, or a
            `"calibrated"` artifact's recorded fingerprint differs from the live model's.
    """
    layout = None
    if session is not None:
        try:
            layout = session.layout
        except Exception:
            layout = None
    if layout is None and model is not None:
        from steerability.algorithms.state_control.common.layout_facts import resolve_layout

        layout = resolve_layout(model, None)
    self._check(layout)
    return self.steering_vector.clone()

_Precomputed

A trivially-resolved source wrapping a concrete SteeringVector (internal).

Lets resolvers treat concrete artifacts and sources uniformly, so precomputed vectors take the same bind path as fitted ones (defensive clone, device/dtype cast). Resolution is model-free, so the source declares ModelAccess.FACTS and runs no fit. Not part of the public API; users pass vectors, mappings, or sources directly.

Source code in steerability/algorithms/state_control/common/sources.py

class _Precomputed:
    """A trivially-resolved source wrapping a concrete `SteeringVector` (internal).

    Lets resolvers treat concrete artifacts and sources uniformly, so precomputed vectors take
    the same bind path as fitted ones (defensive clone, device/dtype cast). Resolution is
    model-free, so the source declares `ModelAccess.FACTS` and runs no fit. Not part of the
    public API; users pass vectors, mappings, or sources directly.
    """

    access: ClassVar[ModelAccess] = ModelAccess.FACTS

    def __init__(self, steering_vector: SteeringVector):
        self._steering_vector = steering_vector

    @property
    def produces_positional(self) -> bool:
        return self._steering_vector.is_positional

    def resolve(
        self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, *, session: "SteeringSession | None" = None
    ) -> SteeringVector:
        return self._steering_vector.clone()

_steering_vector = steering_vector instance-attribute

access = ModelAccess.FACTS class-attribute

produces_positional property

__init__(steering_vector)

Source code in steerability/algorithms/state_control/common/sources.py

def __init__(self, steering_vector: SteeringVector):
    self._steering_vector = steering_vector

resolve(model, tokenizer, *, session=None)

Source code in steerability/algorithms/state_control/common/sources.py

def resolve(
    self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, *, session: "SteeringSession | None" = None
) -> SteeringVector:
    return self._steering_vector.clone()

_as_artifact_source(x)

Coerce a concrete artifact or source into an ArtifactSource (internal).

Accepts an ArtifactSource (returned as-is), a SteeringVector (wrapped in _Precomputed), or a Mapping[int, torch.Tensor] of per-layer directions (wrapped in a SteeringVector with model_type="unknown" then _Precomputed). Anything else raises TypeError.

Source code in steerability/algorithms/state_control/common/sources.py

def _as_artifact_source(x) -> ArtifactSource:
    """Coerce a concrete artifact or source into an `ArtifactSource` (internal).

    Accepts an `ArtifactSource` (returned as-is), a `SteeringVector` (wrapped in `_Precomputed`), or
    a `Mapping[int, torch.Tensor]` of per-layer directions (wrapped in a `SteeringVector` with
    `model_type="unknown"` then `_Precomputed`). Anything else raises `TypeError`.
    """
    if isinstance(x, ArtifactSource):
        return x
    if isinstance(x, SteeringVector):
        return _Precomputed(x)
    if isinstance(x, Mapping):
        directions = {int(k): v for k, v in x.items()}
        if not all(isinstance(v, torch.Tensor) for v in directions.values()):
            raise TypeError("Mapping artifact must map layer ids to torch.Tensor directions.")
        return _Precomputed(SteeringVector(model_type="unknown", directions=directions))
    raise TypeError(
        f"Expected a SteeringVector, a Mapping[int, Tensor], or an ArtifactSource; got "
        f"{type(x).__name__}."
    )

specs

The intervention IR for state control components.

The intervention IR (TokenScope, Intervention) is the single declarative statement of a residual-stream state control's behavior. Both compilers read it: runtime.build_hooks turns a bound intervention tuple into torch hooks for one generation, and lowering.lower_interventions turns it into an InterventionSpec for intervention-capable backends. Components describe their own wire form (WireForm via each component's export), so no layer of the system re-derives another layer's configuration by introspection.

Boundary = Literal['layer_output', 'layer_input'] module-attribute

ScopeKind = Literal['all', 'after_prompt', 'last_k', 'from_position'] module-attribute

Site = Literal['decoder_layer', 'o_proj', 'norm_input'] module-attribute

CoveredLayers

Layer selector resolving to the bound transform's covered layers.

Used when the behavior layers are a fact of the artifact rather than of the model, e.g. a steering plane supplied for a subset of layers. Intervention.bind binds the transform first and takes its covered_layer_ids (intersected with the model's layer range, and with within when given) as the behavior layers, raising when none remain.

Parameters:

Name	Type	Description	Default
within	'Sequence[int] | BaseSelector | None'	Optional base selection the covered layers are intersected with, as explicit layer ids or a selector resolved against the model's layer count.	None

Source code in steerability/algorithms/state_control/common/specs.py

class CoveredLayers:
    """Layer selector resolving to the bound transform's covered layers.

    Used when the behavior layers are a fact of the artifact rather than of the model, e.g. a
    steering plane supplied for a subset of layers. `Intervention.bind` binds the transform
    first and takes its `covered_layer_ids` (intersected with the model's layer range, and
    with `within` when given) as the behavior layers, raising when none remain.

    Args:
        within: Optional base selection the covered layers are intersected with, as explicit
            layer ids or a selector resolved against the model's layer count.
    """

    def __init__(self, within: "Sequence[int] | BaseSelector | None" = None):
        self.within = tuple(int(lid) for lid in within) if isinstance(within, (list, tuple)) else within

    def resolve(self, covered, num_layers: int) -> tuple[int, ...]:
        """The covered layers intersected with the model range and the base selection.

        Raises:
            ValueError: If no layer survives the intersection.
        """
        layer_ids = {int(lid) for lid in covered if 0 <= int(lid) < num_layers}
        if self.within is not None:
            if isinstance(self.within, tuple):
                requested = set(self.within)
            else:
                selected = self.within.select(num_layers=num_layers)
                requested = (
                    {int(lid) for lid in selected}
                    if isinstance(selected, (list, tuple, set, frozenset))
                    else {int(selected)}
                )
            layer_ids &= requested
            if not layer_ids:
                raise ValueError(
                    f"No target layer has a direction in the steering artifact "
                    f"(requested {sorted(requested)}, available {sorted(int(lid) for lid in covered)})."
                )
        if not layer_ids:
            raise ValueError("No active layers for this intervention after filtering.")
        return tuple(sorted(layer_ids))

within = tuple(int(lid) for lid in within) if isinstance(within, (list, tuple)) else within instance-attribute

__init__(within=None)

Source code in steerability/algorithms/state_control/common/specs.py

def __init__(self, within: "Sequence[int] | BaseSelector | None" = None):
    self.within = tuple(int(lid) for lid in within) if isinstance(within, (list, tuple)) else within

resolve(covered, num_layers)

The covered layers intersected with the model range and the base selection.

Raises:

Type	Description
ValueError	If no layer survives the intersection.

Source code in steerability/algorithms/state_control/common/specs.py

def resolve(self, covered, num_layers: int) -> tuple[int, ...]:
    """The covered layers intersected with the model range and the base selection.

    Raises:
        ValueError: If no layer survives the intersection.
    """
    layer_ids = {int(lid) for lid in covered if 0 <= int(lid) < num_layers}
    if self.within is not None:
        if isinstance(self.within, tuple):
            requested = set(self.within)
        else:
            selected = self.within.select(num_layers=num_layers)
            requested = (
                {int(lid) for lid in selected}
                if isinstance(selected, (list, tuple, set, frozenset))
                else {int(selected)}
            )
        layer_ids &= requested
        if not layer_ids:
            raise ValueError(
                f"No target layer has a direction in the steering artifact "
                f"(requested {sorted(requested)}, available {sorted(int(lid) for lid in covered)})."
            )
    if not layer_ids:
        raise ValueError("No active layers for this intervention after filtering.")
    return tuple(sorted(layer_ids))

Intervention dataclass

One activation edit: apply transform at layers, at scope positions, on the boundary side of the layer, whenever gate is open.

Declared unbound at control construction: layers may be a layer selector, the transform may carry an ArtifactSource (or be a factory over a TransformContext), and the gate may be given as a GateSource. bind(model, tokenizer, layout=...) returns the resolved form with layer coverage validated. Kind identity (wire_kinds) is readable on the unbound form, which is what lets check() run before steer().

Interventions are generation-invariant: prompt lengths, pad masks, and position offsets are runtime facts consumed by build_hooks in process and resolved per request by the worker on the wire. Nothing prompt-dependent appears here.

IR dataclasses never use instance defaults for object-valued fields; object-valued defaults use default_factory only.

Attributes:

Name	Type	Description
layers	tuple[int, ...] | 'BaseSelector' | CoveredLayers	Behavior layers (0-based decoder-layer indices), a selector resolved at bind time, or CoveredLayers to take the bound transform's covered layers.
transform	'BaseTransform'	The transform applied at masked positions of open rows.
scope	TokenScope	Token positions to steer.
gate	'Gate | GateSource | None'	The gate consulted at apply time, a source resolving to one, or None for unconditional application.
gate_driven_externally	bool	Follower mode for shared-gate composition. When True, another intervention's condition hooks feed the shared Gate instance, so binding skips the readout compatibility checks and hook compilation builds no condition hooks for this intervention; its behavior hooks only read the shared decision.
boundary	Boundary	Which side of the hooked module the edit applies at. "layer_output" builds forward hooks; "layer_input" builds forward pre-hooks.
site	Site | None	The hooked module family. None derives it from the transform kind (head_additive targets the attention output projection, everything else the decoder layer); "norm_input" targets each layer's normalization sub-modules and has no wire form.
require_coverage	bool	When True (default), bind raises if the resolved transform lacks a direction for any behavior layer; when False, uncovered layers are hooked and pass through unchanged.

Source code in steerability/algorithms/state_control/common/specs.py

@dataclass(frozen=True, slots=True)
class Intervention:
    """One activation edit: apply `transform` at `layers`, at `scope` positions, on the
    `boundary` side of the layer, whenever `gate` is open.

    Declared unbound at control construction: `layers` may be a layer selector, the transform
    may carry an `ArtifactSource` (or be a factory over a `TransformContext`), and the gate may
    be given as a `GateSource`. `bind(model, tokenizer, layout=...)` returns the resolved form
    with layer coverage validated. Kind identity (`wire_kinds`) is readable on the unbound
    form, which is what lets `check()` run before `steer()`.

    Interventions are generation-invariant: prompt lengths, pad masks, and position offsets
    are runtime facts consumed by `build_hooks` in process and resolved per request by the
    worker on the wire. Nothing prompt-dependent appears here.

    IR dataclasses never use instance defaults for object-valued fields; object-valued
    defaults use `default_factory` only.

    Attributes:
        layers: Behavior layers (0-based decoder-layer indices), a selector resolved at bind
            time, or `CoveredLayers` to take the bound transform's covered layers.
        transform: The transform applied at masked positions of open rows.
        scope: Token positions to steer.
        gate: The gate consulted at apply time, a source resolving to one, or None for
            unconditional application.
        gate_driven_externally: Follower mode for shared-gate composition. When True, another
            intervention's condition hooks feed the shared `Gate` instance, so binding skips
            the readout compatibility checks and hook compilation builds no condition hooks
            for this intervention; its behavior hooks only read the shared decision.
        boundary: Which side of the hooked module the edit applies at. `"layer_output"`
            builds forward hooks; `"layer_input"` builds forward pre-hooks.
        site: The hooked module family. None derives it from the transform kind
            (`head_additive` targets the attention output projection, everything else the
            decoder layer); `"norm_input"` targets each layer's normalization sub-modules
            and has no wire form.
        require_coverage: When True (default), `bind` raises if the resolved transform lacks
            a direction for any behavior layer; when False, uncovered layers are hooked and
            pass through unchanged.
    """

    layers: tuple[int, ...] | "BaseSelector" | CoveredLayers
    transform: "BaseTransform"
    scope: TokenScope = field(default_factory=_default_scope)
    gate: "Gate | GateSource | None" = None
    gate_driven_externally: bool = False
    boundary: Boundary = "layer_output"
    site: Site | None = None
    require_coverage: bool = True

    def __post_init__(self):
        if self.boundary not in ("layer_output", "layer_input"):
            raise ValueError(f"boundary must be 'layer_output' or 'layer_input'; got {self.boundary!r}.")
        if self.site not in (None, "decoder_layer", "o_proj", "norm_input"):
            raise ValueError(f"Unknown site {self.site!r}.")
        if isinstance(self.layers, (list, tuple)):
            object.__setattr__(self, "layers", tuple(int(lid) for lid in self.layers))

    @property
    def is_unbound(self) -> bool:
        """True when binding must run model-side work: a layer selector to resolve, a
        transform source or factory to fit, or a gate source to search."""
        from .transforms.base import BaseTransform

        if not isinstance(self.layers, tuple):
            return True
        if not isinstance(self.transform, BaseTransform) or not self.transform.is_bound:
            return True
        if self.gate is not None and not isinstance(self.gate, Gate):
            return True
        return False

    def resolved_site(self) -> Site:
        """The module family this intervention hooks, deriving None from the transform kind."""
        if self.site is not None:
            return self.site
        from .transforms.base import BaseTransform, unwrap_modifiers

        if isinstance(self.transform, BaseTransform):
            core, _ = unwrap_modifiers(self.transform)
            if type(core).wire_kind == "head_additive":
                return "o_proj"
        return "decoder_layer"

    def bind(
        self,
        model: "PreTrainedModel | None",
        tokenizer: "PreTrainedTokenizerBase",
        *,
        layout: "ModelFacts | None" = None,
        session: "SteeringSession | None" = None,
    ) -> "Intervention":
        """Resolve every declared element against `model` (or a session `layout`).

        Resolves the layer selector, binds the transform (fitting artifact sources and
        invoking factories), resolves gate sources, validates layer coverage and readout
        compatibility, and returns the bound intervention. Never mutates `self`.

        Args:
            model: The live model, or None for concrete-artifact configurations bound
                against a session layout.
            tokenizer: Tokenizer used when fitting sources.
            layout: Structural facts (`ModelFacts`) used when `model` is None.
            session: Optional `SteeringSession` forwarded to sources for capture-backed
                fitting and searching.

        Returns:
            The bound intervention.

        Raises:
            ValueError: If a layer is out of range, the transform lacks coverage for a
                behavior layer, the transform's artifact records an extraction location that
                differs from this intervention's boundary, or the gate's readout is
                incompatible with the boundary or model.
        """
        from .layout_facts import resolve_layout
        from .transforms.context import resolve_transform_slot

        layout = layout if layout is not None else resolve_layout(model, session)
        num_layers = layout.num_layers

        transform = self.transform
        if isinstance(self.layers, CoveredLayers):
            transform = resolve_transform_slot(
                transform, model, tokenizer, [], layout=layout,
                require_coverage=False, session=session,
            )
            covered = transform.covered_layer_ids
            if not covered:
                raise ValueError("No active layers for this intervention after filtering.")
            layer_ids = self.layers.resolve(covered, num_layers)
        elif isinstance(self.layers, tuple):
            layer_ids = self.layers
        else:
            selected = self.layers.select(num_layers=num_layers)
            if isinstance(selected, (list, tuple, set, frozenset)):
                layer_ids = tuple(sorted(int(lid) for lid in selected))
            else:
                layer_ids = (int(selected),)
        for lid in layer_ids:
            if not 0 <= lid < num_layers:
                raise ValueError(f"layer_id {lid} out of range [0, {num_layers}).")

        gate = self.gate
        if gate is not None and not isinstance(gate, Gate):
            gate = gate.resolve_gate(model, tokenizer, layout=layout, session=session)
        if gate is not None:
            if not isinstance(gate, Gate):
                raise ValueError(
                    f"gate must resolve to a Gate or None; got {type(gate).__name__}."
                )
            for lid in gate.evidence.layer_ids:
                if not 0 <= lid < num_layers:
                    raise ValueError(f"condition_layer_id {lid} out of range [0, {num_layers}).")
            if not self.gate_driven_externally:
                self._validate_readout(gate.evidence.readout, layout)

        if not isinstance(self.layers, CoveredLayers):
            transform = resolve_transform_slot(
                transform, model, tokenizer, list(layer_ids), layout=layout,
                require_coverage=self.require_coverage, session=session,
            )
        self._validate_transform_artifact(transform)

        bound = replace(self, layers=layer_ids, transform=transform, gate=gate)
        unbound_kinds = self.wire_kinds()
        bound_kinds = bound.wire_kinds()
        # binding may replace parameter values and tensors and may shrink the kind set (a gate
        # source resolving to unconditional drops its readout and rule kinds), never add kinds;
        # narrowing to None is the artifact-dependent case caught by eager steer-time lowering
        assert unbound_kinds is None or bound_kinds is None or unbound_kinds.contains(bound_kinds), (
            f"binding changed wire kinds from {unbound_kinds} to {bound_kinds}"
        )
        return bound

    def _validate_transform_artifact(self, transform) -> None:
        """Check the bound transform's recorded artifact extraction boundary against this
        intervention.

        The check validates recorded provenance only: an artifact whose metadata carries no
        `"location"` key has unknown extraction provenance and passes, with the boundary
        requirement documented on the consuming control's args instead.

        Raises:
            ValueError: If the artifact records a `"location"` that differs from this
                intervention's `boundary`.
        """
        meta = getattr(transform, "artifact_meta", None) or {}
        location = meta.get("location")
        if location is not None and location != self.boundary:
            raise ValueError(
                f"Transform artifact was extracted at '{location}' but this intervention "
                f"hooks '{self.boundary}'. Declare the intervention with "
                f"boundary='{location}', or refit the artifact with "
                f"location='{self.boundary}'."
            )

    def _validate_readout(self, readout, layout) -> None:
        """Check a readout's declared boundary and model identity against this intervention."""
        readout_location = getattr(readout, "location", None)
        if readout_location is not None and readout_location != self.boundary:
            raise ValueError(
                f"Gate readout expects features at '{readout_location}' but this "
                f"intervention hooks '{self.boundary}'. Declare the intervention with "
                f"boundary='{readout_location}', or refit the artifact with "
                f"location='{self.boundary}'."
            )
        readout_fingerprint = getattr(readout, "model_fingerprint", None)
        if readout_fingerprint is not None and layout is not None:
            live_fingerprint = getattr(layout, "model_fingerprint", None)
            if live_fingerprint is not None and readout_fingerprint != live_fingerprint:
                raise ValueError(
                    f"Gate readout was fitted on a different model (fingerprint "
                    f"{readout_fingerprint!r} vs {live_fingerprint!r}). Refit the probe on "
                    "this model, or disarm the check with allow_model_mismatch=True on "
                    "gate_from_probe() or Probe.as_gate()."
                )

    def wire_kinds(self) -> InterventionKinds | None:
        """The wire kind names this configuration lowers to, or None when hook-only.

        Readable on the unbound form: sources and components declare kind identity at
        construction, so `check()` consults this before `steer()`. Artifact-dependent
        inexpressibility (e.g. a resolved artifact missing a direction for an uncovered
        behavior layer) is undetectable here and is caught by eager steer-time lowering.
        """
        from .transforms.base import BaseTransform, unwrap_modifiers

        if self.resolved_site() == "norm_input":
            return None
        modifiers: set[str] = set()
        if isinstance(self.transform, BaseTransform):
            core, wrappers = unwrap_modifiers(self.transform)
            kind = core.wire_plan()
            if kind is None:
                return None
            for wrapper in wrappers:
                modifier_kind = wrapper.modifier_wire_kind(kind)
                if modifier_kind is None:
                    return None
                modifiers.add(modifier_kind)
        else:
            # a factory slot may declare its plan (wire_plan / wire_modifiers); an undeclared
            # factory is unknown before binding
            plan = getattr(self.transform, "wire_plan", None)
            if not callable(plan):
                return None
            kind = plan()
            if kind is None:
                return None
            declared = getattr(self.transform, "wire_modifiers", ())
            modifiers = set(declared() if callable(declared) else declared)
        if (
            self.boundary == "layer_input"
            and self.resolved_site() == "decoder_layer"
            and isinstance(self.layers, tuple)
            and 0 in self.layers
        ):
            # layer 0 input edits precede the first wire boundary; the o_proj site keeps its
            # layer index on the wire, so layer 0 stays expressible there
            return None
        gate = self.gate
        if gate is None:
            readouts: frozenset[str] = frozenset()
            rules: frozenset[str] = frozenset()
        elif isinstance(gate, Gate):
            pair = gate.wire_kinds()
            if pair is None:
                return None
            readouts, rules = pair
        else:
            readouts = getattr(type(gate), "wire_readouts", None)
            rules = getattr(type(gate), "wire_rules", None)
            if readouts is None or rules is None:
                return None
        return InterventionKinds(
            transforms=frozenset({kind}),
            modifiers=frozenset(modifiers),
            scopes=frozenset({self.scope.kind}),
            readouts=readouts,
            rules=rules,
        )

boundary = 'layer_output' class-attribute instance-attribute

gate = None class-attribute instance-attribute

gate_driven_externally = False class-attribute instance-attribute

is_unbound property

True when binding must run model-side work: a layer selector to resolve, a transform source or factory to fit, or a gate source to search.

layers instance-attribute

require_coverage = True class-attribute instance-attribute

scope = field(default_factory=_default_scope) class-attribute instance-attribute

site = None class-attribute instance-attribute

transform instance-attribute

__init__(layers, transform, scope=_default_scope(), gate=None, gate_driven_externally=False, boundary='layer_output', site=None, require_coverage=True)

__post_init__()

Source code in steerability/algorithms/state_control/common/specs.py

def __post_init__(self):
    if self.boundary not in ("layer_output", "layer_input"):
        raise ValueError(f"boundary must be 'layer_output' or 'layer_input'; got {self.boundary!r}.")
    if self.site not in (None, "decoder_layer", "o_proj", "norm_input"):
        raise ValueError(f"Unknown site {self.site!r}.")
    if isinstance(self.layers, (list, tuple)):
        object.__setattr__(self, "layers", tuple(int(lid) for lid in self.layers))

_validate_readout(readout, layout)

Check a readout's declared boundary and model identity against this intervention.

Source code in steerability/algorithms/state_control/common/specs.py

def _validate_readout(self, readout, layout) -> None:
    """Check a readout's declared boundary and model identity against this intervention."""
    readout_location = getattr(readout, "location", None)
    if readout_location is not None and readout_location != self.boundary:
        raise ValueError(
            f"Gate readout expects features at '{readout_location}' but this "
            f"intervention hooks '{self.boundary}'. Declare the intervention with "
            f"boundary='{readout_location}', or refit the artifact with "
            f"location='{self.boundary}'."
        )
    readout_fingerprint = getattr(readout, "model_fingerprint", None)
    if readout_fingerprint is not None and layout is not None:
        live_fingerprint = getattr(layout, "model_fingerprint", None)
        if live_fingerprint is not None and readout_fingerprint != live_fingerprint:
            raise ValueError(
                f"Gate readout was fitted on a different model (fingerprint "
                f"{readout_fingerprint!r} vs {live_fingerprint!r}). Refit the probe on "
                "this model, or disarm the check with allow_model_mismatch=True on "
                "gate_from_probe() or Probe.as_gate()."
            )

_validate_transform_artifact(transform)

Check the bound transform's recorded artifact extraction boundary against this intervention.

The check validates recorded provenance only: an artifact whose metadata carries no "location" key has unknown extraction provenance and passes, with the boundary requirement documented on the consuming control's args instead.

Raises:

Type	Description
ValueError	If the artifact records a "location" that differs from this intervention's boundary.

Source code in steerability/algorithms/state_control/common/specs.py

def _validate_transform_artifact(self, transform) -> None:
    """Check the bound transform's recorded artifact extraction boundary against this
    intervention.

    The check validates recorded provenance only: an artifact whose metadata carries no
    `"location"` key has unknown extraction provenance and passes, with the boundary
    requirement documented on the consuming control's args instead.

    Raises:
        ValueError: If the artifact records a `"location"` that differs from this
            intervention's `boundary`.
    """
    meta = getattr(transform, "artifact_meta", None) or {}
    location = meta.get("location")
    if location is not None and location != self.boundary:
        raise ValueError(
            f"Transform artifact was extracted at '{location}' but this intervention "
            f"hooks '{self.boundary}'. Declare the intervention with "
            f"boundary='{location}', or refit the artifact with "
            f"location='{self.boundary}'."
        )

bind(model, tokenizer, *, layout=None, session=None)

Resolve every declared element against model (or a session layout).

Resolves the layer selector, binds the transform (fitting artifact sources and invoking factories), resolves gate sources, validates layer coverage and readout compatibility, and returns the bound intervention. Never mutates self.

Parameters:

Name	Type	Description	Default
model	'PreTrainedModel | None'	The live model, or None for concrete-artifact configurations bound against a session layout.	required
tokenizer	'PreTrainedTokenizerBase'	Tokenizer used when fitting sources.	required
layout	'ModelFacts | None'	Structural facts (ModelFacts) used when model is None.	None
session	'SteeringSession | None'	Optional SteeringSession forwarded to sources for capture-backed fitting and searching.	None

Returns:

Type	Description
'Intervention'	The bound intervention.

Raises:

Type	Description
ValueError	If a layer is out of range, the transform lacks coverage for a behavior layer, the transform's artifact records an extraction location that differs from this intervention's boundary, or the gate's readout is incompatible with the boundary or model.

Source code in steerability/algorithms/state_control/common/specs.py

def bind(
    self,
    model: "PreTrainedModel | None",
    tokenizer: "PreTrainedTokenizerBase",
    *,
    layout: "ModelFacts | None" = None,
    session: "SteeringSession | None" = None,
) -> "Intervention":
    """Resolve every declared element against `model` (or a session `layout`).

    Resolves the layer selector, binds the transform (fitting artifact sources and
    invoking factories), resolves gate sources, validates layer coverage and readout
    compatibility, and returns the bound intervention. Never mutates `self`.

    Args:
        model: The live model, or None for concrete-artifact configurations bound
            against a session layout.
        tokenizer: Tokenizer used when fitting sources.
        layout: Structural facts (`ModelFacts`) used when `model` is None.
        session: Optional `SteeringSession` forwarded to sources for capture-backed
            fitting and searching.

    Returns:
        The bound intervention.

    Raises:
        ValueError: If a layer is out of range, the transform lacks coverage for a
            behavior layer, the transform's artifact records an extraction location that
            differs from this intervention's boundary, or the gate's readout is
            incompatible with the boundary or model.
    """
    from .layout_facts import resolve_layout
    from .transforms.context import resolve_transform_slot

    layout = layout if layout is not None else resolve_layout(model, session)
    num_layers = layout.num_layers

    transform = self.transform
    if isinstance(self.layers, CoveredLayers):
        transform = resolve_transform_slot(
            transform, model, tokenizer, [], layout=layout,
            require_coverage=False, session=session,
        )
        covered = transform.covered_layer_ids
        if not covered:
            raise ValueError("No active layers for this intervention after filtering.")
        layer_ids = self.layers.resolve(covered, num_layers)
    elif isinstance(self.layers, tuple):
        layer_ids = self.layers
    else:
        selected = self.layers.select(num_layers=num_layers)
        if isinstance(selected, (list, tuple, set, frozenset)):
            layer_ids = tuple(sorted(int(lid) for lid in selected))
        else:
            layer_ids = (int(selected),)
    for lid in layer_ids:
        if not 0 <= lid < num_layers:
            raise ValueError(f"layer_id {lid} out of range [0, {num_layers}).")

    gate = self.gate
    if gate is not None and not isinstance(gate, Gate):
        gate = gate.resolve_gate(model, tokenizer, layout=layout, session=session)
    if gate is not None:
        if not isinstance(gate, Gate):
            raise ValueError(
                f"gate must resolve to a Gate or None; got {type(gate).__name__}."
            )
        for lid in gate.evidence.layer_ids:
            if not 0 <= lid < num_layers:
                raise ValueError(f"condition_layer_id {lid} out of range [0, {num_layers}).")
        if not self.gate_driven_externally:
            self._validate_readout(gate.evidence.readout, layout)

    if not isinstance(self.layers, CoveredLayers):
        transform = resolve_transform_slot(
            transform, model, tokenizer, list(layer_ids), layout=layout,
            require_coverage=self.require_coverage, session=session,
        )
    self._validate_transform_artifact(transform)

    bound = replace(self, layers=layer_ids, transform=transform, gate=gate)
    unbound_kinds = self.wire_kinds()
    bound_kinds = bound.wire_kinds()
    # binding may replace parameter values and tensors and may shrink the kind set (a gate
    # source resolving to unconditional drops its readout and rule kinds), never add kinds;
    # narrowing to None is the artifact-dependent case caught by eager steer-time lowering
    assert unbound_kinds is None or bound_kinds is None or unbound_kinds.contains(bound_kinds), (
        f"binding changed wire kinds from {unbound_kinds} to {bound_kinds}"
    )
    return bound

resolved_site()

The module family this intervention hooks, deriving None from the transform kind.

Source code in steerability/algorithms/state_control/common/specs.py

def resolved_site(self) -> Site:
    """The module family this intervention hooks, deriving None from the transform kind."""
    if self.site is not None:
        return self.site
    from .transforms.base import BaseTransform, unwrap_modifiers

    if isinstance(self.transform, BaseTransform):
        core, _ = unwrap_modifiers(self.transform)
        if type(core).wire_kind == "head_additive":
            return "o_proj"
    return "decoder_layer"

wire_kinds()

The wire kind names this configuration lowers to, or None when hook-only.

Readable on the unbound form: sources and components declare kind identity at construction, so check() consults this before steer(). Artifact-dependent inexpressibility (e.g. a resolved artifact missing a direction for an uncovered behavior layer) is undetectable here and is caught by eager steer-time lowering.

Source code in steerability/algorithms/state_control/common/specs.py

def wire_kinds(self) -> InterventionKinds | None:
    """The wire kind names this configuration lowers to, or None when hook-only.

    Readable on the unbound form: sources and components declare kind identity at
    construction, so `check()` consults this before `steer()`. Artifact-dependent
    inexpressibility (e.g. a resolved artifact missing a direction for an uncovered
    behavior layer) is undetectable here and is caught by eager steer-time lowering.
    """
    from .transforms.base import BaseTransform, unwrap_modifiers

    if self.resolved_site() == "norm_input":
        return None
    modifiers: set[str] = set()
    if isinstance(self.transform, BaseTransform):
        core, wrappers = unwrap_modifiers(self.transform)
        kind = core.wire_plan()
        if kind is None:
            return None
        for wrapper in wrappers:
            modifier_kind = wrapper.modifier_wire_kind(kind)
            if modifier_kind is None:
                return None
            modifiers.add(modifier_kind)
    else:
        # a factory slot may declare its plan (wire_plan / wire_modifiers); an undeclared
        # factory is unknown before binding
        plan = getattr(self.transform, "wire_plan", None)
        if not callable(plan):
            return None
        kind = plan()
        if kind is None:
            return None
        declared = getattr(self.transform, "wire_modifiers", ())
        modifiers = set(declared() if callable(declared) else declared)
    if (
        self.boundary == "layer_input"
        and self.resolved_site() == "decoder_layer"
        and isinstance(self.layers, tuple)
        and 0 in self.layers
    ):
        # layer 0 input edits precede the first wire boundary; the o_proj site keeps its
        # layer index on the wire, so layer 0 stays expressible there
        return None
    gate = self.gate
    if gate is None:
        readouts: frozenset[str] = frozenset()
        rules: frozenset[str] = frozenset()
    elif isinstance(gate, Gate):
        pair = gate.wire_kinds()
        if pair is None:
            return None
        readouts, rules = pair
    else:
        readouts = getattr(type(gate), "wire_readouts", None)
        rules = getattr(type(gate), "wire_rules", None)
        if readouts is None or rules is None:
            return None
    return InterventionKinds(
        transforms=frozenset({kind}),
        modifiers=frozenset(modifiers),
        scopes=frozenset({self.scope.kind}),
        readouts=readouts,
        rules=rules,
    )

TokenScope dataclass

A token-position selector with its parameters.

Attributes:

Name	Type	Description
kind	ScopeKind	One of "all", "after_prompt", "last_k", "from_position".
last_k	int | None	Number of trailing positions, required when kind == "last_k".
from_position	int | None	Absolute start position (inclusive), required when kind == "from_position".

Source code in steerability/algorithms/state_control/common/specs.py

@dataclass(frozen=True, slots=True)
class TokenScope:
    """A token-position selector with its parameters.

    Attributes:
        kind: One of `"all"`, `"after_prompt"`, `"last_k"`, `"from_position"`.
        last_k: Number of trailing positions, required when `kind == "last_k"`.
        from_position: Absolute start position (inclusive), required when
            `kind == "from_position"`.
    """

    kind: ScopeKind
    last_k: int | None = None
    from_position: int | None = None

    def __post_init__(self):
        if self.kind not in get_args(ScopeKind):
            raise ValueError(f"Unknown token scope kind {self.kind!r}.")
        if self.kind == "last_k" and (self.last_k is None or self.last_k < 1):
            raise ValueError("last_k must be >= 1 when kind is 'last_k'.")
        if self.kind == "from_position" and (self.from_position is None or self.from_position < 0):
            raise ValueError("from_position must be >= 0 when kind is 'from_position'.")

    def export(self) -> WireForm:
        """The scope's wire form. Total, since every scope kind is a wire kind."""
        if self.kind == "last_k":
            return WireForm(kind="last_k", params={"k": int(self.last_k)})
        if self.kind == "from_position":
            return WireForm(kind="from_position", params={"position": int(self.from_position)})
        return WireForm(kind=self.kind)

from_position = None class-attribute instance-attribute

kind instance-attribute

last_k = None class-attribute instance-attribute

__init__(kind, last_k=None, from_position=None)

__post_init__()

Source code in steerability/algorithms/state_control/common/specs.py

def __post_init__(self):
    if self.kind not in get_args(ScopeKind):
        raise ValueError(f"Unknown token scope kind {self.kind!r}.")
    if self.kind == "last_k" and (self.last_k is None or self.last_k < 1):
        raise ValueError("last_k must be >= 1 when kind is 'last_k'.")
    if self.kind == "from_position" and (self.from_position is None or self.from_position < 0):
        raise ValueError("from_position must be >= 0 when kind is 'from_position'.")

export()

The scope's wire form. Total, since every scope kind is a wire kind.

Source code in steerability/algorithms/state_control/common/specs.py

def export(self) -> WireForm:
    """The scope's wire form. Total, since every scope kind is a wire kind."""
    if self.kind == "last_k":
        return WireForm(kind="last_k", params={"k": int(self.last_k)})
    if self.kind == "from_position":
        return WireForm(kind="from_position", params={"position": int(self.from_position)})
    return WireForm(kind=self.kind)

WireForm dataclass

One component's form on the wire: the kind name, scalar params, and named tensors.

params follow the plugin's KIND_PARAMS table for the kind; tensors follow its ARTIFACT_TENSORS table.

Attributes:

Name	Type	Description
kind	str	The permanent wire kind name.
params	Mapping[str, float | int | bool | str | tuple[int, ...] | list[int]]	Scalar parameters, inlined next to kind on the wire.
tensors	Mapping[str, Tensor]	Named tensor payloads, materialized as one content-addressed artifact.

Source code in steerability/algorithms/state_control/common/specs.py

@dataclass(frozen=True, slots=True)
class WireForm:
    """One component's form on the wire: the kind name, scalar params, and named tensors.

    `params` follow the plugin's `KIND_PARAMS` table for the kind; `tensors` follow its
    `ARTIFACT_TENSORS` table.

    Attributes:
        kind: The permanent wire kind name.
        params: Scalar parameters, inlined next to `kind` on the wire.
        tensors: Named tensor payloads, materialized as one content-addressed artifact.
    """

    kind: str
    params: Mapping[str, float | int | bool | str | tuple[int, ...] | list[int]] = field(default_factory=dict)
    tensors: Mapping[str, torch.Tensor] = field(default_factory=dict)

kind instance-attribute

params = field(default_factory=dict) class-attribute instance-attribute

tensors = field(default_factory=dict) class-attribute instance-attribute

__init__(kind, params=dict(), tensors=dict())

_default_scope()

Source code in steerability/algorithms/state_control/common/specs.py

def _default_scope() -> TokenScope:
    return TokenScope("after_prompt")

combine_kinds(kind_sets)

Union InterventionKinds across an iterable, propagating None (hook-only).

Source code in steerability/algorithms/state_control/common/specs.py

def combine_kinds(kind_sets) -> InterventionKinds | None:
    """Union `InterventionKinds` across an iterable, propagating None (hook-only)."""
    transforms: set[str] = set()
    modifiers: set[str] = set()
    scopes: set[str] = set()
    readouts: set[str] = set()
    rules: set[str] = set()
    empty = True
    for kinds in kind_sets:
        if kinds is None:
            return None
        empty = False
        transforms |= kinds.transforms
        modifiers |= kinds.modifiers
        scopes |= kinds.scopes
        readouts |= kinds.readouts
        rules |= kinds.rules
    if empty:
        return InterventionKinds()
    return InterventionKinds(
        transforms=frozenset(transforms),
        modifiers=frozenset(modifiers),
        scopes=frozenset(scopes),
        readouts=frozenset(readouts),
        rules=frozenset(rules),
    )

steering_vector

Steering vector: per-layer direction vectors learned by an estimator.

logger = logging.getLogger(__name__) module-attribute

SteeringVector dataclass

Per-layer direction tensors for activation steering.

Directions are stored as [K, D] tensors per layer, where K and D are interpreted by the consuming transform:

CAA, CAST: K=1, D=hidden_size (single broadcast direction)
ActAdd: K=T, D=hidden_size (positional directions)
Angular Steering: K=2, D=hidden_size (orthonormal basis pair)
ITI: K=num_heads, D=head_dim (per-head directions)

The container is agnostic to what K and D mean (varies depending on the method). Semantics come from the consumer (transform).

Attributes:

Name	Type	Description
model_type	str	HuggingFace model_type string (e.g., "llama").
directions	dict[int, Tensor]	Mapping from layer_id to direction tensor of shape [K, D].
num_heads	int | None	Number of attention heads per layer.
head_dim	int | None	Dimension of each head's output.
explained_variances	dict[int, float] | None	Optional mapping from layer_id to explained variance scalar. Only meaningful for estimators that produce a real variance (e.g., PCA-based). None when not applicable.
probe_accuracies	dict[tuple[int, int], float] | None	Optional mapping from (layer_id, head_id) to linear probe validation accuracy (used for head selection in ITI). "Probe" here is ITI's per-head fit-time classifier used to rank heads, distinct from the toolkit's Probe detector.
meta	dict	Provenance record (model, config, tokenizer, and chat-template fingerprints, package version). May be empty for hand-constructed vectors, which disarms cross-backend fingerprint checks.

Source code in steerability/algorithms/state_control/common/steering_vector.py

@dataclass
class SteeringVector:
    """Per-layer direction tensors for activation steering.

    Directions are stored as [K, D] tensors per layer, where K and D are
    interpreted by the consuming transform:

        CAA, CAST: K=1, D=hidden_size (single broadcast direction)
        ActAdd: K=T, D=hidden_size (positional directions)
        Angular Steering: K=2, D=hidden_size (orthonormal basis pair)
        ITI: K=num_heads, D=head_dim (per-head directions)

    The container is agnostic to what K and D mean (varies depending on the method).
    Semantics come from the consumer (transform).

    Attributes:
        model_type: HuggingFace model_type string (e.g., "llama").
        directions: Mapping from layer_id to direction tensor of shape [K, D].
        num_heads: Number of attention heads per layer.
        head_dim: Dimension of each head's output.
        explained_variances: Optional mapping from layer_id to explained
            variance scalar. Only meaningful for estimators that produce a
            real variance (e.g., PCA-based). None when not applicable.
        probe_accuracies: Optional mapping from (layer_id, head_id) to linear
            probe validation accuracy (used for head selection in ITI). "Probe" here is
            ITI's per-head fit-time classifier used to rank heads, distinct from the
            toolkit's `Probe` detector.
        meta: Provenance record (model, config, tokenizer, and chat-template fingerprints,
            package version). May be empty for hand-constructed vectors, which disarms
            cross-backend fingerprint checks.
    """

    model_type: str
    directions: dict[int, torch.Tensor]
    num_heads: int | None = None
    head_dim: int | None = None
    explained_variances: dict[int, float] | None = None
    probe_accuracies: dict[tuple[int, int], float] | None = None
    meta: dict = field(default_factory=dict)

    @property
    def num_tokens(self) -> int:
        """Number of token positions in the steering vector (K dimension)."""
        if not self.directions:
            return 0
        return next(iter(self.directions.values())).size(0)

    @property
    def is_positional(self) -> bool:
        """True if the vector carries per-token positional structure (K > 1)."""
        return self.num_tokens > 1

    def to(self, device: torch.device | str, dtype: torch.dtype | None = None) -> "SteeringVector":
        """Move all direction tensors to device/dtype.

        This mutates `self.directions` in place and returns `self` (so callers may chain).
        To leave the original untouched, `clone()` first.
        """
        self.directions = {
            k: v.to(device=device, dtype=dtype) if dtype else v.to(device=device)
            for k, v in self.directions.items()
        }
        return self

    def clone(self) -> "SteeringVector":
        """Return a deep copy with independent direction tensors and metadata dicts.

        Use this before `to()` / normalization when the vector is caller-supplied and may be
        reused (e.g., shared across controls or across a `ControlSpec` sweep).
        """
        return SteeringVector(
            model_type=self.model_type,
            directions={k: v.clone() for k, v in self.directions.items()},
            num_heads=self.num_heads,
            head_dim=self.head_dim,
            explained_variances=dict(self.explained_variances) if self.explained_variances is not None else None,
            probe_accuracies=dict(self.probe_accuracies) if self.probe_accuracies is not None else None,
            meta=dict(self.meta),
        )

    def normalized(self) -> "SteeringVector":
        """Return a per-layer L2-normalized deep clone.

        Each direction tensor is normalized by its own L2 norm (over all elements). Rows with a
        zero (or near-zero) norm are left unchanged. This is the bound-mode twin of
        `ContrastiveFit(normalize=True)`: the original vector is untouched, so a shared/caller-held
        vector stays intact.

        Returns:
            A new `SteeringVector` with normalized directions and copied metadata.
        """
        clone = self.clone()
        for layer_id, direction in clone.directions.items():
            norm = direction.norm()
            if norm > 0:
                clone.directions[layer_id] = direction / norm
        return clone

    def scaled_to_norms(
        self,
        target_norms: Mapping[int, float],
        scale: float = 1.0,
    ) -> "SteeringVector":
        """Return a deep clone whose per-layer directions are rescaled to a target L2 norm.

        At each layer present in both `target_norms` and `self.directions`, the direction is rescaled
        to L2 norm `scale * target_norms[layer]` while keeping its orientation, scaling each direction
        to a caller-chosen (typically residual-norm-relative) magnitude. The original vector is
        untouched.

        Semantics are defined for broadcast directions only (`K == 1`), where every steered
        token in the layer receives the same vector.

        Args:
            target_norms: Mapping from layer id to the per-layer base norm to scale to (e.g. the
                output of `measure_residual_norms`). Layers absent from either `target_norms` or
                `self.directions` are dropped from the result.
            scale: Multiplier applied to each target norm (the "dose"). Must be positive.

        Returns:
            A new `SteeringVector` covering the intersection of `target_norms` and `self.directions`,
            each direction stored as `[1, H]`, with metadata copied from the original.

        Raises:
            ValueError: If `scale <= 0`, any covered target norm is non-positive, a covered source
                direction is positional (`K > 1`) or has zero norm, or the layer intersection is
                empty.
        """
        if scale <= 0:
            raise ValueError(f"scale must be positive, got {scale}.")

        covered = [lid for lid in target_norms if lid in self.directions]
        if not covered:
            raise ValueError(
                "scaled_to_norms: no overlap between target_norms layers "
                f"{sorted(target_norms.keys())} and vector layers {sorted(self.directions.keys())}."
            )

        clone = self.clone()
        scaled: dict[int, torch.Tensor] = {}
        for lid in covered:
            target = float(target_norms[lid])
            if target <= 0:
                raise ValueError(f"target norm for layer {lid} must be positive, got {target}.")

            direction = clone.directions[lid]
            if direction.ndim == 2 and direction.shape[0] == 1:
                direction = direction.squeeze(0)
            elif direction.ndim != 1:
                raise ValueError(
                    f"scaled_to_norms is defined for broadcast directions (K=1); layer {lid} has "
                    f"shape {tuple(clone.directions[lid].shape)}."
                )

            source_norm = direction.norm()
            if source_norm <= 0:
                raise ValueError(
                    f"layer {lid} has a zero-norm source direction; a target norm is unreachable "
                    "from the zero vector."
                )

            scaled[lid] = (direction * (scale * target / source_norm)).unsqueeze(0)  # [1, H]

        clone.directions = scaled
        return clone

    def validate(self) -> None:
        """Validate that required fields are populated.

        Raises:
            ValueError: If model_type or directions are empty.
        """
        if not self.model_type:
            raise ValueError("model_type must be provided.")
        if not self.directions:
            raise ValueError("directions must not be empty.")

    def save(self, file_path: str) -> None:
        """Save the SteeringVector to a JSON file.

        Args:
            file_path: Path to save to. ".svec" extension added if not present.
        """
        if not file_path.endswith(".svec"):
            file_path += ".svec"
        directory = os.path.dirname(file_path)

        if directory:
            os.makedirs(directory, exist_ok=True)
        data: dict = {
            "model_type": self.model_type,
            "directions": {str(k): v.tolist() for k, v in self.directions.items()},
        }

        if self.num_heads is not None:
            data["num_heads"] = self.num_heads
        if self.head_dim is not None:
            data["head_dim"] = self.head_dim
        if self.explained_variances is not None:
            data["explained_variances"] = {str(k): v for k, v in self.explained_variances.items()}
        if self.probe_accuracies is not None:
            data["probe_accuracies"] = {
                f"{layer}:{head}": acc for (layer, head), acc in self.probe_accuracies.items()
            }
        if self.meta:
            data["meta"] = self.meta

        with open(file_path, "w") as f:
            json.dump(data, f)
        logger.debug("Saved SteeringVector to %s", file_path)

    @classmethod
    def load(cls, file_path: str) -> "SteeringVector":
        """Load a SteeringVector from a JSON file.

        Args:
            file_path: Path to load from. ".svec" extension added if not present.

        Returns:
            Loaded SteeringVector instance.
        """
        if not file_path.endswith(".svec"):
            file_path += ".svec"
        with open(file_path) as f:
            data = json.load(f)

        directions = {}
        for k, v in data["directions"].items():
            t = torch.tensor(v, dtype=torch.float32)
            if t.ndim == 1:
                t = t.unsqueeze(0)  # [D] -> [1, D] backward compatibility
            directions[int(k)] = t

        explained_variances = None
        if "explained_variances" in data:
            explained_variances = {int(k): float(v) for k, v in data["explained_variances"].items()}

        num_heads = data.get("num_heads")
        head_dim = data.get("head_dim")

        probe_accuracies = None
        if "probe_accuracies" in data:
            probe_accuracies = {}
            for k, acc in data["probe_accuracies"].items():
                layer_str, head_str = k.split(":")
                probe_accuracies[(int(layer_str), int(head_str))] = float(acc)

        logger.debug("Loaded SteeringVector from %s with layers %s", file_path, list(directions.keys()))
        return cls(
            model_type=data["model_type"],
            directions=directions,
            num_heads=num_heads,
            head_dim=head_dim,
            explained_variances=explained_variances,
            probe_accuracies=probe_accuracies,
            meta=data.get("meta", {}),
        )

directions instance-attribute

explained_variances = None class-attribute instance-attribute

head_dim = None class-attribute instance-attribute

is_positional property

True if the vector carries per-token positional structure (K > 1).

meta = field(default_factory=dict) class-attribute instance-attribute

model_type instance-attribute

num_heads = None class-attribute instance-attribute

num_tokens property

Number of token positions in the steering vector (K dimension).

probe_accuracies = None class-attribute instance-attribute

__init__(model_type, directions, num_heads=None, head_dim=None, explained_variances=None, probe_accuracies=None, meta=dict())

clone()

Return a deep copy with independent direction tensors and metadata dicts.

Use this before to() / normalization when the vector is caller-supplied and may be reused (e.g., shared across controls or across a ControlSpec sweep).

Source code in steerability/algorithms/state_control/common/steering_vector.py

def clone(self) -> "SteeringVector":
    """Return a deep copy with independent direction tensors and metadata dicts.

    Use this before `to()` / normalization when the vector is caller-supplied and may be
    reused (e.g., shared across controls or across a `ControlSpec` sweep).
    """
    return SteeringVector(
        model_type=self.model_type,
        directions={k: v.clone() for k, v in self.directions.items()},
        num_heads=self.num_heads,
        head_dim=self.head_dim,
        explained_variances=dict(self.explained_variances) if self.explained_variances is not None else None,
        probe_accuracies=dict(self.probe_accuracies) if self.probe_accuracies is not None else None,
        meta=dict(self.meta),
    )

load(file_path) classmethod

Load a SteeringVector from a JSON file.

Parameters:

Name	Type	Description	Default
file_path	str	Path to load from. ".svec" extension added if not present.	required

Returns:

Type	Description
SteeringVector	Loaded SteeringVector instance.

Source code in steerability/algorithms/state_control/common/steering_vector.py

@classmethod
def load(cls, file_path: str) -> "SteeringVector":
    """Load a SteeringVector from a JSON file.

    Args:
        file_path: Path to load from. ".svec" extension added if not present.

    Returns:
        Loaded SteeringVector instance.
    """
    if not file_path.endswith(".svec"):
        file_path += ".svec"
    with open(file_path) as f:
        data = json.load(f)

    directions = {}
    for k, v in data["directions"].items():
        t = torch.tensor(v, dtype=torch.float32)
        if t.ndim == 1:
            t = t.unsqueeze(0)  # [D] -> [1, D] backward compatibility
        directions[int(k)] = t

    explained_variances = None
    if "explained_variances" in data:
        explained_variances = {int(k): float(v) for k, v in data["explained_variances"].items()}

    num_heads = data.get("num_heads")
    head_dim = data.get("head_dim")

    probe_accuracies = None
    if "probe_accuracies" in data:
        probe_accuracies = {}
        for k, acc in data["probe_accuracies"].items():
            layer_str, head_str = k.split(":")
            probe_accuracies[(int(layer_str), int(head_str))] = float(acc)

    logger.debug("Loaded SteeringVector from %s with layers %s", file_path, list(directions.keys()))
    return cls(
        model_type=data["model_type"],
        directions=directions,
        num_heads=num_heads,
        head_dim=head_dim,
        explained_variances=explained_variances,
        probe_accuracies=probe_accuracies,
        meta=data.get("meta", {}),
    )

normalized()

Return a per-layer L2-normalized deep clone.

Each direction tensor is normalized by its own L2 norm (over all elements). Rows with a zero (or near-zero) norm are left unchanged. This is the bound-mode twin of ContrastiveFit(normalize=True): the original vector is untouched, so a shared/caller-held vector stays intact.

Returns:

Type	Description
SteeringVector	A new SteeringVector with normalized directions and copied metadata.

Source code in steerability/algorithms/state_control/common/steering_vector.py

def normalized(self) -> "SteeringVector":
    """Return a per-layer L2-normalized deep clone.

    Each direction tensor is normalized by its own L2 norm (over all elements). Rows with a
    zero (or near-zero) norm are left unchanged. This is the bound-mode twin of
    `ContrastiveFit(normalize=True)`: the original vector is untouched, so a shared/caller-held
    vector stays intact.

    Returns:
        A new `SteeringVector` with normalized directions and copied metadata.
    """
    clone = self.clone()
    for layer_id, direction in clone.directions.items():
        norm = direction.norm()
        if norm > 0:
            clone.directions[layer_id] = direction / norm
    return clone

save(file_path)

Save the SteeringVector to a JSON file.

Parameters:

Name	Type	Description	Default
file_path	str	Path to save to. ".svec" extension added if not present.	required

Source code in steerability/algorithms/state_control/common/steering_vector.py

def save(self, file_path: str) -> None:
    """Save the SteeringVector to a JSON file.

    Args:
        file_path: Path to save to. ".svec" extension added if not present.
    """
    if not file_path.endswith(".svec"):
        file_path += ".svec"
    directory = os.path.dirname(file_path)

    if directory:
        os.makedirs(directory, exist_ok=True)
    data: dict = {
        "model_type": self.model_type,
        "directions": {str(k): v.tolist() for k, v in self.directions.items()},
    }

    if self.num_heads is not None:
        data["num_heads"] = self.num_heads
    if self.head_dim is not None:
        data["head_dim"] = self.head_dim
    if self.explained_variances is not None:
        data["explained_variances"] = {str(k): v for k, v in self.explained_variances.items()}
    if self.probe_accuracies is not None:
        data["probe_accuracies"] = {
            f"{layer}:{head}": acc for (layer, head), acc in self.probe_accuracies.items()
        }
    if self.meta:
        data["meta"] = self.meta

    with open(file_path, "w") as f:
        json.dump(data, f)
    logger.debug("Saved SteeringVector to %s", file_path)

scaled_to_norms(target_norms, scale=1.0)

Return a deep clone whose per-layer directions are rescaled to a target L2 norm.

At each layer present in both target_norms and self.directions, the direction is rescaled to L2 norm scale * target_norms[layer] while keeping its orientation, scaling each direction to a caller-chosen (typically residual-norm-relative) magnitude. The original vector is untouched.

Semantics are defined for broadcast directions only (K == 1), where every steered token in the layer receives the same vector.

Parameters:

Name	Type	Description	Default
target_norms	Mapping[int, float]	Mapping from layer id to the per-layer base norm to scale to (e.g. the output of measure_residual_norms). Layers absent from either target_norms or self.directions are dropped from the result.	required
scale	float	Multiplier applied to each target norm (the "dose"). Must be positive.	1.0

Returns:

Type	Description
SteeringVector	A new SteeringVector covering the intersection of target_norms and self.directions,
SteeringVector	each direction stored as [1, H], with metadata copied from the original.

Raises:

Type	Description
ValueError	If scale <= 0, any covered target norm is non-positive, a covered source direction is positional (K > 1) or has zero norm, or the layer intersection is empty.

Source code in steerability/algorithms/state_control/common/steering_vector.py

def scaled_to_norms(
    self,
    target_norms: Mapping[int, float],
    scale: float = 1.0,
) -> "SteeringVector":
    """Return a deep clone whose per-layer directions are rescaled to a target L2 norm.

    At each layer present in both `target_norms` and `self.directions`, the direction is rescaled
    to L2 norm `scale * target_norms[layer]` while keeping its orientation, scaling each direction
    to a caller-chosen (typically residual-norm-relative) magnitude. The original vector is
    untouched.

    Semantics are defined for broadcast directions only (`K == 1`), where every steered
    token in the layer receives the same vector.

    Args:
        target_norms: Mapping from layer id to the per-layer base norm to scale to (e.g. the
            output of `measure_residual_norms`). Layers absent from either `target_norms` or
            `self.directions` are dropped from the result.
        scale: Multiplier applied to each target norm (the "dose"). Must be positive.

    Returns:
        A new `SteeringVector` covering the intersection of `target_norms` and `self.directions`,
        each direction stored as `[1, H]`, with metadata copied from the original.

    Raises:
        ValueError: If `scale <= 0`, any covered target norm is non-positive, a covered source
            direction is positional (`K > 1`) or has zero norm, or the layer intersection is
            empty.
    """
    if scale <= 0:
        raise ValueError(f"scale must be positive, got {scale}.")

    covered = [lid for lid in target_norms if lid in self.directions]
    if not covered:
        raise ValueError(
            "scaled_to_norms: no overlap between target_norms layers "
            f"{sorted(target_norms.keys())} and vector layers {sorted(self.directions.keys())}."
        )

    clone = self.clone()
    scaled: dict[int, torch.Tensor] = {}
    for lid in covered:
        target = float(target_norms[lid])
        if target <= 0:
            raise ValueError(f"target norm for layer {lid} must be positive, got {target}.")

        direction = clone.directions[lid]
        if direction.ndim == 2 and direction.shape[0] == 1:
            direction = direction.squeeze(0)
        elif direction.ndim != 1:
            raise ValueError(
                f"scaled_to_norms is defined for broadcast directions (K=1); layer {lid} has "
                f"shape {tuple(clone.directions[lid].shape)}."
            )

        source_norm = direction.norm()
        if source_norm <= 0:
            raise ValueError(
                f"layer {lid} has a zero-norm source direction; a target norm is unreachable "
                "from the zero vector."
            )

        scaled[lid] = (direction * (scale * target / source_norm)).unsqueeze(0)  # [1, H]

    clone.directions = scaled
    return clone

to(device, dtype=None)

Move all direction tensors to device/dtype.

This mutates self.directions in place and returns self (so callers may chain). To leave the original untouched, clone() first.

Source code in steerability/algorithms/state_control/common/steering_vector.py

def to(self, device: torch.device | str, dtype: torch.dtype | None = None) -> "SteeringVector":
    """Move all direction tensors to device/dtype.

    This mutates `self.directions` in place and returns `self` (so callers may chain).
    To leave the original untouched, `clone()` first.
    """
    self.directions = {
        k: v.to(device=device, dtype=dtype) if dtype else v.to(device=device)
        for k, v in self.directions.items()
    }
    return self

validate()

Validate that required fields are populated.

Raises:

Type	Description
ValueError	If model_type or directions are empty.

Source code in steerability/algorithms/state_control/common/steering_vector.py

def validate(self) -> None:
    """Validate that required fields are populated.

    Raises:
        ValueError: If model_type or directions are empty.
    """
    if not self.model_type:
        raise ValueError("model_type must be provided.")
    if not self.directions:
        raise ValueError("directions must not be empty.")

token_scope

Token scope utilities for computing position masks.

align_mask_to_batch(mask, hidden_batch)

Expand a [B, T] mask to the hidden-state batch dimension.

Masks are built with B = prompt batch size, but HuggingFace generate expands hidden states to B * num_beams via repeat_interleave for beam search. Row order is therefore [item0, item0, ..., item1, item1, ...]; this replicates the mask the same way so it aligns when a transform broadcasts it against the expanded hidden states.

Parameters:

Name	Type	Description	Default
mask	BoolTensor	Boolean mask of shape [B, T].	required
hidden_batch	int	The batch size of the hidden states the mask will be applied to.	required

Returns:

Type	Description
BoolTensor	A [hidden_batch, T] mask (the input unchanged when hidden_batch == B).

Raises:

Type	Description
RuntimeError	If hidden_batch is not a multiple of B (unexpected generation expansion).

Source code in steerability/algorithms/state_control/common/token_scope.py

def align_mask_to_batch(mask: torch.BoolTensor, hidden_batch: int) -> torch.BoolTensor:
    """Expand a `[B, T]` mask to the hidden-state batch dimension.

    Masks are built with `B` = prompt batch size, but HuggingFace `generate` expands hidden
    states to `B * num_beams` via `repeat_interleave` for beam search. Row order is therefore
    `[item0, item0, ..., item1, item1, ...]`; this replicates the mask the same way so it aligns
    when a transform broadcasts it against the expanded hidden states.

    Args:
        mask: Boolean mask of shape `[B, T]`.
        hidden_batch: The batch size of the hidden states the mask will be applied to.

    Returns:
        A `[hidden_batch, T]` mask (the input unchanged when `hidden_batch == B`).

    Raises:
        RuntimeError: If `hidden_batch` is not a multiple of `B` (unexpected generation expansion).
    """
    B = mask.size(0)
    if hidden_batch == B:
        return mask
    if hidden_batch % B == 0:
        return mask.repeat_interleave(hidden_batch // B, dim=0)
    raise RuntimeError(
        f"Hidden batch {hidden_batch} is not a multiple of the prompt batch {B}; "
        f"cannot align steering mask (unexpected generation expansion)."
    )

compute_prompt_lens(input_ids, pad_token_id=None)

Compute per-batch-item prompt lengths from input_ids.

For the "after_prompt" token scope, the prompt length should represent the absolute position where generation begins. With KV-cached generation, this is always the full input sequence length (seq_len), regardless of padding side or pad token count. The model's KV cache covers all input positions including pads, and generation continues from position seq_len.

Parameters:

Name	Type	Description	Default
input_ids	LongTensor	Shape [B, T] or [T].	required
pad_token_id	int | None	Unused. Kept for API compatibility.	None

Returns:

Type	Description
LongTensor	Tensor of shape [B] where each value is seq_len.

Source code in steerability/algorithms/state_control/common/token_scope.py

def compute_prompt_lens(
    input_ids: torch.LongTensor,
    pad_token_id: int | None = None,  # noqa: ARG001
) -> torch.LongTensor:
    """Compute per-batch-item prompt lengths from input_ids.

    For the "after_prompt" token scope, the prompt length should represent the
    absolute position where generation begins. With KV-cached generation, this
    is always the full input sequence length (seq_len), regardless of padding
    side or pad token count. The model's KV cache covers all input positions
    including pads, and generation continues from position seq_len.

    Args:
        input_ids: Shape [B, T] or [T].
        pad_token_id: Unused. Kept for API compatibility.

    Returns:
        Tensor of shape [B] where each value is seq_len.
    """
    if input_ids.ndim == 1:
        input_ids = input_ids.unsqueeze(0)
    return torch.full(
        (input_ids.size(0),),
        input_ids.size(1),
        dtype=torch.long,
        device=input_ids.device,
    )

make_token_mask(scope, *, seq_len, prompt_lens, last_k=None, from_position=None, position_offset=0)

Build a [B, T] boolean mask selecting which tokens to transform.

Parameters:

Name	Type	Description	Default
scope	ScopeKind	Which tokens to include. "all" - every position is True. "after_prompt" - only positions beyond the prompt are True. During autoregressive generation with KV cache, the hook may see only newly generated tokens (seq_len=1). Use position_offset to indicate the cumulative position in the full sequence. "last_k" - only the last k positions are True. "from_position" - positions >= from_position are True. Useful for single forward pass evaluation (e.g., logit scoring) where you want to steer from a specific point within the prompt.	required
seq_len	int	Current sequence length T (may be 1 during KV-cached generation).	required
prompt_lens	LongTensor	Shape [B], per-item prompt lengths.	required
last_k	int | None	Required when scope == "last_k".	None
from_position	int | None	Required when scope == "from_position". The absolute position from which to start steering (inclusive).	None
position_offset	int	Cumulative position offset for KV-cached generation. When the model processes token N in the sequence but only passes a single token to the hook (seq_len=1), set position_offset=N so that "after_prompt" correctly identifies generated tokens.	0

Returns:

Type	Description
BoolTensor	Boolean tensor of shape [B, T].

Source code in steerability/algorithms/state_control/common/token_scope.py

def make_token_mask(
    scope: ScopeKind,
    *,
    seq_len: int,
    prompt_lens: torch.LongTensor,
    last_k: int | None = None,
    from_position: int | None = None,
    position_offset: int = 0,
) -> torch.BoolTensor:
    """Build a [B, T] boolean mask selecting which tokens to transform.

    Args:
        scope: Which tokens to include.
            "all" - every position is True.
            "after_prompt" - only positions beyond the prompt are True.
                During autoregressive generation with KV cache, the hook may
                see only newly generated tokens (seq_len=1). Use position_offset
                to indicate the cumulative position in the full sequence.
            "last_k" - only the last k positions are True.
            "from_position" - positions >= from_position are True. Useful for
                single forward pass evaluation (e.g., logit scoring) where you
                want to steer from a specific point within the prompt.
        seq_len: Current sequence length T (may be 1 during KV-cached generation).
        prompt_lens: Shape [B], per-item prompt lengths.
        last_k: Required when scope == "last_k".
        from_position: Required when scope == "from_position". The absolute
            position from which to start steering (inclusive).
        position_offset: Cumulative position offset for KV-cached generation.
            When the model processes token N in the sequence but only passes
            a single token to the hook (seq_len=1), set position_offset=N so
            that "after_prompt" correctly identifies generated tokens.

    Returns:
        Boolean tensor of shape [B, T].
    """
    B = prompt_lens.size(0)
    device = prompt_lens.device

    # compute absolute positions in the full sequence
    local_positions = torch.arange(seq_len, device=device).unsqueeze(0).expand(B, -1)
    absolute_positions = local_positions + position_offset

    if scope == "all":
        return torch.ones(B, seq_len, dtype=torch.bool, device=device)
    elif scope == "after_prompt":
        return absolute_positions >= prompt_lens.unsqueeze(1)
    elif scope == "last_k":
        if last_k is None or last_k < 1:
            raise ValueError("last_k must be >= 1 when scope is 'last_k'.")
        # for last_k, we want the last k positions relative to current seq_len
        return local_positions >= (seq_len - last_k)
    elif scope == "from_position":
        if from_position is None or from_position < 0:
            raise ValueError("from_position must be >= 0 when scope is 'from_position'.")
        return absolute_positions >= from_position
    else:
        raise ValueError(f"Unknown token scope: {scope!r}")

transforms

Transform components for state control.

additive

Additive activation steering transform.

AdditiveTransform

Bases: BaseTransform

Adds scaled direction vector(s) to hidden states.

Two modes, selected by the positional flag:

Broadcast (default, e.g., CAA): h'[pos] = h[pos] + mask[pos] * strength * direction[0]

The same `[1, H]` vector is added at every masked position. The ``alignment``
parameter is unused.

Positional (e.g., ActAdd): h'[p] = h[p] + mask[p] * strength * direction[p - alignment]

Row `t` of a `[T, H]` direction is added at absolute sequence position
``alignment + t``. The window ``[alignment, alignment + T)`` is evaluated in
absolute positions, so during KV-cached generation a window inside the prompt
fires only on the prefill pass, and a window extending past the prompt fires at
exactly the covered generated positions, once each.

Direction tensors are canonicalized to [T, H] at construction or bind (a 1-D [H] tensor becomes [1, H]). A direction with T > 1 requires positional=True, since broadcast semantics are undefined for multi-row directions.

Parameters:

Name	Type	Description	Default
artifact	SteeringVector | Mapping[int, Tensor] | ArtifactSource	The steering artifact, given as a SteeringVector, a per-layer directions mapping (Mapping[int, Tensor], each [T, H] or [H]), or an ArtifactSource (unbound until bind(ctx)). Required.	required
strength	float	Global scaling factor.	1.0
alignment	int	Absolute token position of the first direction row (default: 0). Used only when positional is True.	0
positional	bool	Selects per-position injection. When False (default), the direction broadcasts over masked positions and every layer's direction must have T == 1.	False

Raises:

Type	Description
ValueError	If a concrete direction has T > 1 while positional is False (raised at construction for concrete artifacts, at bind for sources).
TypeError	If artifact is not a supported type.

Source code in steerability/algorithms/state_control/common/transforms/additive.py

class AdditiveTransform(BaseTransform):
    """Adds scaled direction vector(s) to hidden states.

    Two modes, selected by the `positional` flag:

    Broadcast (default, e.g., CAA):
        ``h'[pos] = h[pos] + mask[pos] * strength * direction[0]``

        The same `[1, H]` vector is added at every masked position. The ``alignment``
        parameter is unused.

    Positional (e.g., ActAdd):
        ``h'[p] = h[p] + mask[p] * strength * direction[p - alignment]``

        Row `t` of a `[T, H]` direction is added at absolute sequence position
        ``alignment + t``. The window ``[alignment, alignment + T)`` is evaluated in
        absolute positions, so during KV-cached generation a window inside the prompt
        fires only on the prefill pass, and a window extending past the prompt fires at
        exactly the covered generated positions, once each.

    Direction tensors are canonicalized to `[T, H]` at construction or bind (a 1-D `[H]`
    tensor becomes `[1, H]`). A direction with `T > 1` requires `positional=True`, since
    broadcast semantics are undefined for multi-row directions.

    Args:
        artifact: The steering artifact, given as a `SteeringVector`, a per-layer directions
            mapping (`Mapping[int, Tensor]`, each `[T, H]` or `[H]`), or an `ArtifactSource`
            (unbound until `bind(ctx)`). Required.
        strength: Global scaling factor.
        alignment: Absolute token position of the first direction row (default: 0). Used
            only when `positional` is True.
        positional: Selects per-position injection. When False (default), the direction
            broadcasts over masked positions and every layer's direction must have `T == 1`.

    Raises:
        ValueError: If a concrete direction has `T > 1` while `positional` is False (raised
            at construction for concrete artifacts, at `bind` for sources).
        TypeError: If `artifact` is not a supported type.
    """

    wire_kind: ClassVar[str | None] = "additive"

    def __init__(
        self,
        artifact: SteeringVector | Mapping[int, torch.Tensor] | ArtifactSource,
        strength: float = 1.0,
        alignment: int = 0,
        positional: bool = False,
    ):
        self.strength = strength
        self.alignment = alignment
        self.positional = positional
        self._source: ArtifactSource | None = None
        self.directions: dict[int, torch.Tensor] | None = None

        self._artifact_meta: dict | None = None
        if isinstance(artifact, ArtifactSource):
            self._source = artifact
        elif isinstance(artifact, SteeringVector):
            self.directions = _canonical_directions(artifact.directions, positional)
            self._artifact_meta = dict(artifact.meta) if artifact.meta else None
        elif isinstance(artifact, Mapping):
            self.directions = _canonical_directions(artifact, positional)
        else:
            raise TypeError(
                f"AdditiveTransform artifact must be a SteeringVector, a Mapping[int, Tensor], or an "
                f"ArtifactSource; got {type(artifact).__name__} (did you mean strength=?)."
            )

    @property
    def is_bound(self) -> bool:
        return self.directions is not None

    @property
    def artifact_meta(self) -> dict | None:
        return self._artifact_meta

    def bind(self, ctx: "TransformContext") -> "AdditiveTransform":
        if self.is_bound:
            return self
        return AdditiveTransform(
            ctx.resolve(self._source),
            strength=self.strength,
            alignment=self.alignment,
            positional=self.positional,
        )

    @property
    def covered_layer_ids(self) -> set[int] | None:
        return set(self.directions.keys()) if self.directions is not None else None

    def to_config(self) -> tuple[dict, object | None, "BaseTransform | None"]:
        """The `(params, artifact, inner)` serialized form under the `additive` kind.

        The artifact slot carries the concrete directions as a `SteeringVector` when bound,
        or the unresolved source when not.
        """
        params = {"strength": float(self.strength), "alignment": int(self.alignment),
                  "positional": bool(self.positional)}
        if self.directions is not None:
            artifact = SteeringVector(
                model_type="unknown", directions=dict(self.directions), meta=dict(self._artifact_meta or {}),
            )
        else:
            artifact = self._source
        return params, artifact, None

    @classmethod
    def from_config(cls, params: dict, *, artifact=None, inner=None) -> "AdditiveTransform":
        """Rebuild an `additive` transform from its serialized form."""
        return cls(
            artifact,
            strength=params.get("strength", 1.0),
            alignment=params.get("alignment", 0),
            positional=params.get("positional", False),
        )

    def wire_plan(self) -> str | None:
        """`"additive"` for broadcast transforms; None when `positional` is True.

        Positional injection has no wire form regardless of `T`.
        """
        return None if self.positional else "additive"

    def export(self, layer_id: int) -> "WireForm | None":
        """The `additive` wire form for `layer_id`, or None for positional transforms.

        Semantics are defined for broadcast directions only, where every steered token
        receives the same vector.
        """
        from ..specs import WireForm

        if self.positional or self.directions is None:
            return None
        direction = self.directions.get(layer_id)
        if direction is None:
            return None
        return WireForm(
            kind="additive",
            params={"strength": float(self.strength)},
            tensors={"vector": direction.squeeze(0)},
        )


    def apply(
        self,
        hidden_states: torch.Tensor,
        *,
        layer_id: int,
        token_mask: torch.BoolTensor,
        position_offset: int = 0,
        **kwargs,
    ) -> torch.Tensor:
        """Apply additive steering.

        Args:
            hidden_states: Shape [B, T_seq, H].
            layer_id: Which layer this is being applied at.
            token_mask: Shape [B, T_seq]. True at positions to modify.
            position_offset: Absolute position of the pass's first token. Consumed in
                positional mode, where the local slice covering absolute positions
                `[alignment, alignment + T)` receives the direction rows; ignored in
                broadcast mode.

        Returns:
            Modified hidden states, same shape as input.
        """
        self._require_bound()
        direction = self.directions.get(layer_id)
        if direction is None:
            return hidden_states

        if not self.positional:
            # broadcast mode (e.g., CAA); same vector at all masked positions
            v = (self.strength * direction.squeeze(0)).to(
                dtype=hidden_states.dtype,
                device=hidden_states.device,
            )
            delta = token_mask.unsqueeze(-1).to(hidden_states.dtype) * v.view(1, 1, -1)
            return hidden_states + delta

        # positional mode (e.g., ActAdd); the pass covers absolute positions
        # [position_offset, position_offset + seq_len), the window covers [a, a + T)
        seq_len = hidden_states.size(1)
        T_steer = direction.size(0)
        a = self.alignment

        window_start = max(a, position_offset)
        window_end = min(a + T_steer, position_offset + seq_len)
        if window_start >= window_end:
            return hidden_states

        local_start = window_start - position_offset
        local_end = window_end - position_offset
        vec_start = window_start - a
        vec_end = vec_start + (window_end - window_start)

        v = (self.strength * direction[vec_start:vec_end]).to(
            dtype=hidden_states.dtype,
            device=hidden_states.device,
        )  # [inject_len, H]

        mask_slice = token_mask[:, local_start:local_end]  # [B, inject_len]
        gated_v = mask_slice.unsqueeze(-1).to(hidden_states.dtype) * v.unsqueeze(0)

        # add in-place at the injection slice
        out = hidden_states.clone()
        out[:, local_start:local_end] += gated_v
        return out

_artifact_meta = None instance-attribute

_source = None instance-attribute

alignment = alignment instance-attribute

artifact_meta property

Provenance metadata of the transform's steering artifact, or None.

Populated when the artifact was supplied or resolved as a SteeringVector carrying meta (fit fingerprints); consumers cross-check it against a serving engine's model identity. Wrappers delegate to their inner transform. The default returns None.

covered_layer_ids property

Layers this transform can act on; None = unknown (opts out of adapter validation).

Artifact transforms report set(directions.keys()) when bound and None when unbound; wrappers delegate to their inner transform.

directions = None instance-attribute

is_bound property

True when the transform can be applied as-is (artifact concrete).

Default True: transforms that take no steering artifact are always bound.

is_modifier = False class-attribute

positional = positional instance-attribute

source property

The unresolved ArtifactSource this transform carries, or None when concrete.

strength = strength instance-attribute

wire_kind = 'additive' class-attribute

__init__(artifact, strength=1.0, alignment=0, positional=False)

Source code in steerability/algorithms/state_control/common/transforms/additive.py

def __init__(
    self,
    artifact: SteeringVector | Mapping[int, torch.Tensor] | ArtifactSource,
    strength: float = 1.0,
    alignment: int = 0,
    positional: bool = False,
):
    self.strength = strength
    self.alignment = alignment
    self.positional = positional
    self._source: ArtifactSource | None = None
    self.directions: dict[int, torch.Tensor] | None = None

    self._artifact_meta: dict | None = None
    if isinstance(artifact, ArtifactSource):
        self._source = artifact
    elif isinstance(artifact, SteeringVector):
        self.directions = _canonical_directions(artifact.directions, positional)
        self._artifact_meta = dict(artifact.meta) if artifact.meta else None
    elif isinstance(artifact, Mapping):
        self.directions = _canonical_directions(artifact, positional)
    else:
        raise TypeError(
            f"AdditiveTransform artifact must be a SteeringVector, a Mapping[int, Tensor], or an "
            f"ArtifactSource; got {type(artifact).__name__} (did you mean strength=?)."
        )

_require_bound()

Raise a clear RuntimeError if this transform is unbound.

Called as the first line of every artifact-consuming apply.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def _require_bound(self) -> None:
    """Raise a clear `RuntimeError` if this transform is unbound.

    Called as the first line of every artifact-consuming `apply`.
    """
    if not self.is_bound:
        raise RuntimeError(
            f"{type(self).__name__} was constructed with an ArtifactSource and is unbound; pass "
            f"it through a control that resolves and binds it during steer() (ActivationAdapter, "
            f"or CAST via behavior_transform)."
        )

apply(hidden_states, *, layer_id, token_mask, position_offset=0, **kwargs)

Apply additive steering.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Shape [B, T_seq, H].	required
layer_id	int	Which layer this is being applied at.	required
token_mask	BoolTensor	Shape [B, T_seq]. True at positions to modify.	required
position_offset	int	Absolute position of the pass's first token. Consumed in positional mode, where the local slice covering absolute positions [alignment, alignment + T) receives the direction rows; ignored in broadcast mode.	0

Returns:

Type	Description
Tensor	Modified hidden states, same shape as input.

Source code in steerability/algorithms/state_control/common/transforms/additive.py

def apply(
    self,
    hidden_states: torch.Tensor,
    *,
    layer_id: int,
    token_mask: torch.BoolTensor,
    position_offset: int = 0,
    **kwargs,
) -> torch.Tensor:
    """Apply additive steering.

    Args:
        hidden_states: Shape [B, T_seq, H].
        layer_id: Which layer this is being applied at.
        token_mask: Shape [B, T_seq]. True at positions to modify.
        position_offset: Absolute position of the pass's first token. Consumed in
            positional mode, where the local slice covering absolute positions
            `[alignment, alignment + T)` receives the direction rows; ignored in
            broadcast mode.

    Returns:
        Modified hidden states, same shape as input.
    """
    self._require_bound()
    direction = self.directions.get(layer_id)
    if direction is None:
        return hidden_states

    if not self.positional:
        # broadcast mode (e.g., CAA); same vector at all masked positions
        v = (self.strength * direction.squeeze(0)).to(
            dtype=hidden_states.dtype,
            device=hidden_states.device,
        )
        delta = token_mask.unsqueeze(-1).to(hidden_states.dtype) * v.view(1, 1, -1)
        return hidden_states + delta

    # positional mode (e.g., ActAdd); the pass covers absolute positions
    # [position_offset, position_offset + seq_len), the window covers [a, a + T)
    seq_len = hidden_states.size(1)
    T_steer = direction.size(0)
    a = self.alignment

    window_start = max(a, position_offset)
    window_end = min(a + T_steer, position_offset + seq_len)
    if window_start >= window_end:
        return hidden_states

    local_start = window_start - position_offset
    local_end = window_end - position_offset
    vec_start = window_start - a
    vec_end = vec_start + (window_end - window_start)

    v = (self.strength * direction[vec_start:vec_end]).to(
        dtype=hidden_states.dtype,
        device=hidden_states.device,
    )  # [inject_len, H]

    mask_slice = token_mask[:, local_start:local_end]  # [B, inject_len]
    gated_v = mask_slice.unsqueeze(-1).to(hidden_states.dtype) * v.unsqueeze(0)

    # add in-place at the injection slice
    out = hidden_states.clone()
    out[:, local_start:local_end] += gated_v
    return out

bind(ctx)

Return a fully-bound transform for this context.

Contract:

- MUST NOT mutate `self`; instances and sources are shared across adapters and
  `ControlSpec` grid points, whose params objects are reused per point.
- Returns `self` when already bound (idempotent).
- When source-carrying, returns a NEW instance of the same class constructed with
  `ctx.resolve(self._source)` and all hyperparameters copied, built by fresh
  construction rather than `copy.copy`.

Default: returns self.

Source code in steerability/algorithms/state_control/common/transforms/additive.py

def bind(self, ctx: "TransformContext") -> "AdditiveTransform":
    if self.is_bound:
        return self
    return AdditiveTransform(
        ctx.resolve(self._source),
        strength=self.strength,
        alignment=self.alignment,
        positional=self.positional,
    )

export(layer_id)

The additive wire form for layer_id, or None for positional transforms.

Semantics are defined for broadcast directions only, where every steered token receives the same vector.

Source code in steerability/algorithms/state_control/common/transforms/additive.py

def export(self, layer_id: int) -> "WireForm | None":
    """The `additive` wire form for `layer_id`, or None for positional transforms.

    Semantics are defined for broadcast directions only, where every steered token
    receives the same vector.
    """
    from ..specs import WireForm

    if self.positional or self.directions is None:
        return None
    direction = self.directions.get(layer_id)
    if direction is None:
        return None
    return WireForm(
        kind="additive",
        params={"strength": float(self.strength)},
        tensors={"vector": direction.squeeze(0)},
    )

export_modifier(layer_id)

The wrapper's wire modifier form at layer_id, or None to contribute no modifier at that layer.

Meaningful only on wrapper transforms (is_modifier True); kind-level inexpressibility is reported by modifier_wire_kind. The default returns None.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def export_modifier(self, layer_id: int) -> "WireForm | None":
    """The wrapper's wire modifier form at `layer_id`, or None to contribute no modifier
    at that layer.

    Meaningful only on wrapper transforms (`is_modifier` True); kind-level
    inexpressibility is reported by `modifier_wire_kind`. The default returns None.
    """
    return None

from_config(params, *, artifact=None, inner=None) classmethod

Rebuild an additive transform from its serialized form.

Source code in steerability/algorithms/state_control/common/transforms/additive.py

@classmethod
def from_config(cls, params: dict, *, artifact=None, inner=None) -> "AdditiveTransform":
    """Rebuild an `additive` transform from its serialized form."""
    return cls(
        artifact,
        strength=params.get("strength", 1.0),
        alignment=params.get("alignment", 0),
        positional=params.get("positional", False),
    )

modifier_wire_kind(core_kind)

The wire modifier kind this wrapper contributes over a core transform kind, or None when the combination has no wire form.

Meaningful only on wrapper transforms (is_modifier True). The default returns the class wire_kind.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def modifier_wire_kind(self, core_kind: str) -> str | None:
    """The wire modifier kind this wrapper contributes over a core transform kind, or
    None when the combination has no wire form.

    Meaningful only on wrapper transforms (`is_modifier` True). The default returns the
    class `wire_kind`.
    """
    return type(self).wire_kind

to_config()

The (params, artifact, inner) serialized form under the additive kind.

The artifact slot carries the concrete directions as a SteeringVector when bound, or the unresolved source when not.

Source code in steerability/algorithms/state_control/common/transforms/additive.py

def to_config(self) -> tuple[dict, object | None, "BaseTransform | None"]:
    """The `(params, artifact, inner)` serialized form under the `additive` kind.

    The artifact slot carries the concrete directions as a `SteeringVector` when bound,
    or the unresolved source when not.
    """
    params = {"strength": float(self.strength), "alignment": int(self.alignment),
              "positional": bool(self.positional)}
    if self.directions is not None:
        artifact = SteeringVector(
            model_type="unknown", directions=dict(self.directions), meta=dict(self._artifact_meta or {}),
        )
    else:
        artifact = self._source
    return params, artifact, None

wire_plan()

"additive" for broadcast transforms; None when positional is True.

Positional injection has no wire form regardless of T.

Source code in steerability/algorithms/state_control/common/transforms/additive.py

def wire_plan(self) -> str | None:
    """`"additive"` for broadcast transforms; None when `positional` is True.

    Positional injection has no wire form regardless of `T`.
    """
    return None if self.positional else "additive"

_canonical_directions(directions, positional)

Canonicalize per-layer directions to [T, H], validating shape against the mode.

A 1-D [H] tensor becomes [1, H]. Broadcast mode (positional=False) requires T == 1 at every layer.

Raises:

Type	Description
ValueError	If a direction is not 1-D or 2-D, or has T > 1 while positional is False.

Source code in steerability/algorithms/state_control/common/transforms/additive.py

def _canonical_directions(directions: Mapping[int, torch.Tensor], positional: bool) -> dict[int, torch.Tensor]:
    """Canonicalize per-layer directions to `[T, H]`, validating shape against the mode.

    A 1-D `[H]` tensor becomes `[1, H]`. Broadcast mode (`positional=False`) requires
    `T == 1` at every layer.

    Raises:
        ValueError: If a direction is not 1-D or 2-D, or has `T > 1` while `positional` is
            False.
    """
    canonical: dict[int, torch.Tensor] = {}
    for layer_id, direction in directions.items():
        if direction.ndim == 1:
            direction = direction.unsqueeze(0)
        if direction.ndim != 2:
            raise ValueError(
                f"AdditiveTransform direction for layer {layer_id} must be [H] or [T, H]; got "
                f"shape {tuple(direction.shape)}."
            )
        if not positional and direction.size(0) != 1:
            raise ValueError(
                f"AdditiveTransform direction for layer {layer_id} has T={direction.size(0)} rows "
                f"but positional=False; broadcast semantics are defined for single-row directions "
                f"only. Construct the transform with positional=True for per-position injection."
            )
        canonical[int(layer_id)] = direction
    return canonical

alignment_adaptive

Alignment-adaptive filtering: a transform decorator that gates by feature alignment.

AlignmentAdaptiveTransform

Bases: BaseTransform

Applies an inner transform only at tokens aligned with a feature direction.

Computes each token's projection onto a chosen direction and keeps only positions whose alignment exceeds threshold; those at or below are removed from the token mask before the inner transform runs. With threshold=0.0 this applies the Adaptive Angular Steering gate mask = max(0, sign(h·d_feat)) (Eq. 3 of Vu and Nguyen), which rotates only tokens already positively aligned with the feature axis. Because it only narrows the mask, it composes with any BaseTransform.

The transform carries two artifacts, the inner transform's and its own alignment axis, and is bound iff both are bound. bind recurses into the inner and resolves its own source; the two may share the same source instance, in which case memoization runs the shared fit once.

Parameters:

Name	Type	Description	Default
inner	BaseTransform	The transform to apply at the surviving positions.	required
artifact	SteeringVector | Mapping[int, Tensor] | ArtifactSource	The steering artifact whose per-layer directions supply the alignment axis, given as a SteeringVector, a per-layer directions mapping, or an ArtifactSource (unbound until bind(ctx)). The axis used is directions[layer_id][direction_index]. Required.	required
threshold	float	Alignment cutoff; positions with alignment > threshold survive.	0.0
direction_index	int	Row of the per-layer direction tensor to use as the alignment axis (row 0 = feature axis for angular steering).	0
use_cosine	bool	If True, alignment is cosine similarity; otherwise it is the projection onto the unit-normalized direction.	False

Reference:

• "Angular Steering: Behavior Control via Rotation in Activation Space" Hieu M. Vu, Tan M. Nguyen https://arxiv.org/abs/2510.26243

Source code in steerability/algorithms/state_control/common/transforms/alignment_adaptive.py

class AlignmentAdaptiveTransform(BaseTransform):
    """Applies an inner transform only at tokens aligned with a feature direction.

    Computes each token's projection onto a chosen direction and keeps only positions whose
    alignment exceeds `threshold`; those at or below are removed from the token mask before the
    inner transform runs. With `threshold=0.0` this applies the Adaptive Angular Steering gate
    `mask = max(0, sign(h·d_feat))` (Eq. 3 of Vu and Nguyen), which rotates only tokens already
    positively aligned with the feature axis. Because it only narrows the mask, it composes with
    any `BaseTransform`.

    The transform carries two artifacts, the inner transform's and its own alignment axis, and is
    bound iff both are bound. `bind` recurses into the inner and resolves its own source; the two
    may share the same source instance, in which case memoization runs the shared fit once.

    Args:
        inner: The transform to apply at the surviving positions.
        artifact: The steering artifact whose per-layer directions supply the alignment axis, given
            as a `SteeringVector`, a per-layer directions mapping, or an `ArtifactSource` (unbound
            until `bind(ctx)`). The axis used is `directions[layer_id][direction_index]`. Required.
        threshold: Alignment cutoff; positions with alignment `> threshold` survive.
        direction_index: Row of the per-layer direction tensor to use as the alignment axis
            (row 0 = feature axis for angular steering).
        use_cosine: If True, alignment is cosine similarity; otherwise it is the projection onto
            the unit-normalized direction.

    Reference:

    - "Angular Steering: Behavior Control via Rotation in Activation Space"
      Hieu M. Vu, Tan M. Nguyen
      [https://arxiv.org/abs/2510.26243](https://arxiv.org/abs/2510.26243)
    """

    wire_kind: ClassVar[str | None] = "alignment_adaptive"
    is_modifier: ClassVar[bool] = True

    def __init__(
        self,
        inner: BaseTransform,
        artifact: SteeringVector | Mapping[int, torch.Tensor] | ArtifactSource,
        threshold: float = 0.0,
        direction_index: int = 0,
        use_cosine: bool = False,
    ):
        self.inner = inner
        self.threshold = float(threshold)
        self.direction_index = int(direction_index)
        self.use_cosine = bool(use_cosine)
        self._source: ArtifactSource | None = None
        self.steering_vector: SteeringVector | None = None

        if isinstance(artifact, ArtifactSource):
            self._source = artifact
        elif isinstance(artifact, SteeringVector):
            self.steering_vector = artifact
        elif isinstance(artifact, Mapping):
            self.steering_vector = SteeringVector(model_type="unknown", directions=dict(artifact))
        else:
            raise TypeError(
                f"AlignmentAdaptiveTransform artifact must be a SteeringVector, a Mapping[int, Tensor], "
                f"or an ArtifactSource; got {type(artifact).__name__} (did you mean threshold=?)."
            )

    @property
    def is_bound(self) -> bool:
        return self.steering_vector is not None and self.inner.is_bound

    @property
    def artifact_meta(self) -> dict | None:
        return self.inner.artifact_meta

    def bind(self, ctx: "TransformContext") -> "AlignmentAdaptiveTransform":
        if self.is_bound:
            return self
        own = self.steering_vector if self.steering_vector is not None else ctx.resolve(self._source)
        return AlignmentAdaptiveTransform(
            self.inner.bind(ctx),
            own,
            threshold=self.threshold,
            direction_index=self.direction_index,
            use_cosine=self.use_cosine,
        )

    @property
    def covered_layer_ids(self) -> set[int] | None:
        return self.inner.covered_layer_ids

    def to_config(self) -> tuple[dict, object | None, BaseTransform | None]:
        """The `(params, artifact, inner)` serialized form under the `alignment_adaptive` kind."""
        params = {
            "threshold": float(self.threshold),
            "direction_index": int(self.direction_index),
            "use_cosine": bool(self.use_cosine),
        }
        artifact = self.steering_vector if self.steering_vector is not None else self._source
        return params, artifact, self.inner

    @classmethod
    def from_config(cls, params: dict, *, artifact=None, inner=None) -> "AlignmentAdaptiveTransform":
        """Rebuild an `alignment_adaptive` wrapper from its serialized form."""
        return cls(
            inner,
            artifact,
            threshold=params.get("threshold", 0.0),
            direction_index=params.get("direction_index", 0),
            use_cosine=params.get("use_cosine", False),
        )

    def modifier_wire_kind(self, core_kind: str) -> str | None:
        """`"alignment_adaptive"`, or None over a per-head core.

        Like `norm_preserving`, the wire modifier operates on the residual stream; a wrapped
        `head_additive` is hook-only.
        """
        if core_kind == "head_additive":
            return None
        return "alignment_adaptive"

    def export_modifier(self, layer_id: int) -> "WireForm | None":
        """The `alignment_adaptive` wire modifier form at `layer_id`.

        The modifier's wire vector is the resolved per-layer alignment axis
        (`directions[layer_id][direction_index]`). A layer without an alignment axis
        contributes no modifier, matching the in-process behavior where the mask is left
        unnarrowed there.
        """
        from ..specs import WireForm

        if self.steering_vector is None:
            return None
        dirs = self.steering_vector.directions.get(layer_id)
        if dirs is None:
            return None
        return WireForm(
            kind="alignment_adaptive",
            params={"threshold": float(self.threshold), "use_cosine": bool(self.use_cosine)},
            tensors={"vector": dirs[self.direction_index]},
        )


    def apply(
        self,
        hidden_states: torch.Tensor,
        *,
        layer_id: int,
        token_mask: torch.BoolTensor,
        **kwargs,
    ) -> torch.Tensor:
        """Narrow the token mask by feature alignment, then delegate to the inner transform.

        Args:
            hidden_states: Shape `[B, T, H]`.
            layer_id: Which layer this is being applied at.
            token_mask: Shape `[B, T]`. True at positions the inner transform may modify.

        Returns:
            Modified hidden states, same shape as input.
        """
        self._require_bound()
        dirs = self.steering_vector.directions.get(layer_id)
        if dirs is not None:
            direction = dirs[self.direction_index].to(hidden_states.device, hidden_states.dtype)
            if self.use_cosine:
                alignment = torch.nn.functional.cosine_similarity(
                    hidden_states, direction.view(1, 1, -1), dim=-1
                )
            else:
                unit = direction / (direction.norm() + 1e-8)
                alignment = hidden_states @ unit  # [B, T]
            token_mask = token_mask & (alignment > self.threshold)
        return self.inner.apply(hidden_states, layer_id=layer_id, token_mask=token_mask, **kwargs)

_source = None instance-attribute

artifact_meta property

Provenance metadata of the transform's steering artifact, or None.

Populated when the artifact was supplied or resolved as a SteeringVector carrying meta (fit fingerprints); consumers cross-check it against a serving engine's model identity. Wrappers delegate to their inner transform. The default returns None.

covered_layer_ids property

Layers this transform can act on; None = unknown (opts out of adapter validation).

Artifact transforms report set(directions.keys()) when bound and None when unbound; wrappers delegate to their inner transform.

direction_index = int(direction_index) instance-attribute

inner = inner instance-attribute

is_bound property

True when the transform can be applied as-is (artifact concrete).

Default True: transforms that take no steering artifact are always bound.

is_modifier = True class-attribute

source property

The unresolved ArtifactSource this transform carries, or None when concrete.

steering_vector = None instance-attribute

threshold = float(threshold) instance-attribute

use_cosine = bool(use_cosine) instance-attribute

wire_kind = 'alignment_adaptive' class-attribute

__init__(inner, artifact, threshold=0.0, direction_index=0, use_cosine=False)

Source code in steerability/algorithms/state_control/common/transforms/alignment_adaptive.py

def __init__(
    self,
    inner: BaseTransform,
    artifact: SteeringVector | Mapping[int, torch.Tensor] | ArtifactSource,
    threshold: float = 0.0,
    direction_index: int = 0,
    use_cosine: bool = False,
):
    self.inner = inner
    self.threshold = float(threshold)
    self.direction_index = int(direction_index)
    self.use_cosine = bool(use_cosine)
    self._source: ArtifactSource | None = None
    self.steering_vector: SteeringVector | None = None

    if isinstance(artifact, ArtifactSource):
        self._source = artifact
    elif isinstance(artifact, SteeringVector):
        self.steering_vector = artifact
    elif isinstance(artifact, Mapping):
        self.steering_vector = SteeringVector(model_type="unknown", directions=dict(artifact))
    else:
        raise TypeError(
            f"AlignmentAdaptiveTransform artifact must be a SteeringVector, a Mapping[int, Tensor], "
            f"or an ArtifactSource; got {type(artifact).__name__} (did you mean threshold=?)."
        )

_require_bound()

Raise a clear RuntimeError if this transform is unbound.

Called as the first line of every artifact-consuming apply.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def _require_bound(self) -> None:
    """Raise a clear `RuntimeError` if this transform is unbound.

    Called as the first line of every artifact-consuming `apply`.
    """
    if not self.is_bound:
        raise RuntimeError(
            f"{type(self).__name__} was constructed with an ArtifactSource and is unbound; pass "
            f"it through a control that resolves and binds it during steer() (ActivationAdapter, "
            f"or CAST via behavior_transform)."
        )

apply(hidden_states, *, layer_id, token_mask, **kwargs)

Narrow the token mask by feature alignment, then delegate to the inner transform.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Shape [B, T, H].	required
layer_id	int	Which layer this is being applied at.	required
token_mask	BoolTensor	Shape [B, T]. True at positions the inner transform may modify.	required

Returns:

Type	Description
Tensor	Modified hidden states, same shape as input.

Source code in steerability/algorithms/state_control/common/transforms/alignment_adaptive.py

def apply(
    self,
    hidden_states: torch.Tensor,
    *,
    layer_id: int,
    token_mask: torch.BoolTensor,
    **kwargs,
) -> torch.Tensor:
    """Narrow the token mask by feature alignment, then delegate to the inner transform.

    Args:
        hidden_states: Shape `[B, T, H]`.
        layer_id: Which layer this is being applied at.
        token_mask: Shape `[B, T]`. True at positions the inner transform may modify.

    Returns:
        Modified hidden states, same shape as input.
    """
    self._require_bound()
    dirs = self.steering_vector.directions.get(layer_id)
    if dirs is not None:
        direction = dirs[self.direction_index].to(hidden_states.device, hidden_states.dtype)
        if self.use_cosine:
            alignment = torch.nn.functional.cosine_similarity(
                hidden_states, direction.view(1, 1, -1), dim=-1
            )
        else:
            unit = direction / (direction.norm() + 1e-8)
            alignment = hidden_states @ unit  # [B, T]
        token_mask = token_mask & (alignment > self.threshold)
    return self.inner.apply(hidden_states, layer_id=layer_id, token_mask=token_mask, **kwargs)

bind(ctx)

Return a fully-bound transform for this context.

Contract:

- MUST NOT mutate `self`; instances and sources are shared across adapters and
  `ControlSpec` grid points, whose params objects are reused per point.
- Returns `self` when already bound (idempotent).
- When source-carrying, returns a NEW instance of the same class constructed with
  `ctx.resolve(self._source)` and all hyperparameters copied, built by fresh
  construction rather than `copy.copy`.

Default: returns self.

Source code in steerability/algorithms/state_control/common/transforms/alignment_adaptive.py

def bind(self, ctx: "TransformContext") -> "AlignmentAdaptiveTransform":
    if self.is_bound:
        return self
    own = self.steering_vector if self.steering_vector is not None else ctx.resolve(self._source)
    return AlignmentAdaptiveTransform(
        self.inner.bind(ctx),
        own,
        threshold=self.threshold,
        direction_index=self.direction_index,
        use_cosine=self.use_cosine,
    )

export(layer_id)

This configuration's wire form at layer_id, or None when the configuration is not expressible in the wire vocabulary.

Exportability is a property of a configuration, not a class; an artifact whose shape has no wire semantics (e.g. a positional direction) returns None even though the class names a wire_kind. The default returns None.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def export(self, layer_id: int) -> "WireForm | None":
    """This configuration's wire form at `layer_id`, or None when the configuration is
    not expressible in the wire vocabulary.

    Exportability is a property of a configuration, not a class; an artifact whose shape
    has no wire semantics (e.g. a positional direction) returns None even though the
    class names a `wire_kind`. The default returns None.
    """
    return None

export_modifier(layer_id)

The alignment_adaptive wire modifier form at layer_id.

The modifier's wire vector is the resolved per-layer alignment axis (directions[layer_id][direction_index]). A layer without an alignment axis contributes no modifier, matching the in-process behavior where the mask is left unnarrowed there.

Source code in steerability/algorithms/state_control/common/transforms/alignment_adaptive.py

def export_modifier(self, layer_id: int) -> "WireForm | None":
    """The `alignment_adaptive` wire modifier form at `layer_id`.

    The modifier's wire vector is the resolved per-layer alignment axis
    (`directions[layer_id][direction_index]`). A layer without an alignment axis
    contributes no modifier, matching the in-process behavior where the mask is left
    unnarrowed there.
    """
    from ..specs import WireForm

    if self.steering_vector is None:
        return None
    dirs = self.steering_vector.directions.get(layer_id)
    if dirs is None:
        return None
    return WireForm(
        kind="alignment_adaptive",
        params={"threshold": float(self.threshold), "use_cosine": bool(self.use_cosine)},
        tensors={"vector": dirs[self.direction_index]},
    )

from_config(params, *, artifact=None, inner=None) classmethod

Rebuild an alignment_adaptive wrapper from its serialized form.

Source code in steerability/algorithms/state_control/common/transforms/alignment_adaptive.py

@classmethod
def from_config(cls, params: dict, *, artifact=None, inner=None) -> "AlignmentAdaptiveTransform":
    """Rebuild an `alignment_adaptive` wrapper from its serialized form."""
    return cls(
        inner,
        artifact,
        threshold=params.get("threshold", 0.0),
        direction_index=params.get("direction_index", 0),
        use_cosine=params.get("use_cosine", False),
    )

modifier_wire_kind(core_kind)

"alignment_adaptive", or None over a per-head core.

Like norm_preserving, the wire modifier operates on the residual stream; a wrapped head_additive is hook-only.

Source code in steerability/algorithms/state_control/common/transforms/alignment_adaptive.py

def modifier_wire_kind(self, core_kind: str) -> str | None:
    """`"alignment_adaptive"`, or None over a per-head core.

    Like `norm_preserving`, the wire modifier operates on the residual stream; a wrapped
    `head_additive` is hook-only.
    """
    if core_kind == "head_additive":
        return None
    return "alignment_adaptive"

to_config()

The (params, artifact, inner) serialized form under the alignment_adaptive kind.

Source code in steerability/algorithms/state_control/common/transforms/alignment_adaptive.py

def to_config(self) -> tuple[dict, object | None, BaseTransform | None]:
    """The `(params, artifact, inner)` serialized form under the `alignment_adaptive` kind."""
    params = {
        "threshold": float(self.threshold),
        "direction_index": int(self.direction_index),
        "use_cosine": bool(self.use_cosine),
    }
    artifact = self.steering_vector if self.steering_vector is not None else self._source
    return params, artifact, self.inner

wire_plan()

The wire transform kind this configuration serializes to, or None when hook-only.

Readable on the unbound form: a source-carrying transform consults its source's declared shape rather than resolving it. The default returns the class wire_kind.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def wire_plan(self) -> str | None:
    """The wire transform kind this configuration serializes to, or None when hook-only.

    Readable on the unbound form: a source-carrying transform consults its source's
    declared shape rather than resolving it. The default returns the class `wire_kind`.
    """
    return type(self).wire_kind

base

Base class for hidden-state transforms.

BaseTransform

Bases: ABC

Applies a modification to hidden states at a given layer.

All transforms receive

• hidden_states shaped [B, T, H]
• the layer_id so the transform can index per-layer artifacts
• a token_mask shaped [B, T] (True where the transform should apply)

Transforms MUST NOT modify hidden_states in-place if the original tensor is needed later (e.g., for norm-preserving wrappers); return a new tensor.

Artifact-carrying transforms take their artifact as the first positional argument, typed SteeringVector | Mapping[int, Tensor] | ArtifactSource and required. A concrete artifact binds the transform immediately (validated at __init__); an ArtifactSource leaves it unbound until bind(ctx) resolves the source (validated then). Subclass recipe:

- store the concrete artifact and set `is_bound=True`, or store the source and set
  `is_bound=False`;
- override `bind(ctx)` to return a freshly constructed bound instance, never mutating
  `self` and never using `copy.copy`, so that derived caches stay valid;
- override `covered_layer_ids` to report the layers the (bound) artifact covers;
- call `self._require_bound()` as the first line of `apply`.

Transforms with no steering artifact keep the defaults (always bound, no coverage).

Class attributes

wire_kind: The permanent wire kind name this class serializes to, or None when the class has no wire form. Wire names mirror toolkit class names, so the mapping is definitional rather than maintained. is_modifier: True for wrapper transforms that hold an inner transform and serialize as a wire modifier rather than a transform kind.

Source code in steerability/algorithms/state_control/common/transforms/base.py

class BaseTransform(ABC):
    """Applies a modification to hidden states at a given layer.

    All transforms receive:
        - hidden_states shaped [B, T, H]
        - the layer_id so the transform can index per-layer artifacts
        - a token_mask shaped [B, T] (True where the transform should apply)

    Transforms MUST NOT modify hidden_states in-place if the original tensor
    is needed later (e.g., for norm-preserving wrappers); return a new tensor.

    Artifact-carrying transforms take their artifact as the first positional argument, typed
    `SteeringVector | Mapping[int, Tensor] | ArtifactSource` and required. A concrete artifact
    binds the transform immediately (validated at `__init__`); an `ArtifactSource` leaves it
    unbound until `bind(ctx)` resolves the source (validated then). Subclass recipe:

        - store the concrete artifact and set `is_bound=True`, or store the source and set
          `is_bound=False`;
        - override `bind(ctx)` to return a freshly constructed bound instance, never mutating
          `self` and never using `copy.copy`, so that derived caches stay valid;
        - override `covered_layer_ids` to report the layers the (bound) artifact covers;
        - call `self._require_bound()` as the first line of `apply`.

    Transforms with no steering artifact keep the defaults (always bound, no coverage).

    Class attributes:
        wire_kind: The permanent wire kind name this class serializes to, or None when the
            class has no wire form. Wire names mirror toolkit class names, so the mapping is
            definitional rather than maintained.
        is_modifier: True for wrapper transforms that hold an `inner` transform and serialize
            as a wire modifier rather than a transform kind.
    """

    wire_kind: ClassVar[str | None] = None
    is_modifier: ClassVar[bool] = False

    @property
    def artifact_meta(self) -> dict | None:
        """Provenance metadata of the transform's steering artifact, or None.

        Populated when the artifact was supplied or resolved as a `SteeringVector` carrying
        `meta` (fit fingerprints); consumers cross-check it against a serving engine's model
        identity. Wrappers delegate to their inner transform. The default returns None.
        """
        return None

    @abstractmethod
    def apply(
        self,
        hidden_states: torch.Tensor,
        *,
        layer_id: int,
        token_mask: torch.BoolTensor,
        **kwargs: Any,
    ) -> torch.Tensor:
        """Apply the transform and return modified hidden states.

        Args:
            hidden_states: Shape [B, T, H].
            layer_id: Which layer this is being applied at.
            token_mask: Shape [B, T]. True at positions to modify.
            **kwargs: Transform-specific extra arguments.

        Returns:
            Modified hidden states, same shape as input.
        """
        ...

    @property
    def is_bound(self) -> bool:
        """True when the transform can be applied as-is (artifact concrete).

        Default True: transforms that take no steering artifact are always bound.
        """
        return True

    @property
    def source(self):
        """The unresolved `ArtifactSource` this transform carries, or None when concrete."""
        return getattr(self, "_source", None)

    def bind(self, ctx: "TransformContext") -> "BaseTransform":
        """Return a fully-bound transform for this context.

        Contract:

            - MUST NOT mutate `self`; instances and sources are shared across adapters and
              `ControlSpec` grid points, whose params objects are reused per point.
            - Returns `self` when already bound (idempotent).
            - When source-carrying, returns a NEW instance of the same class constructed with
              `ctx.resolve(self._source)` and all hyperparameters copied, built by fresh
              construction rather than `copy.copy`.

        Default: returns self.
        """
        return self

    def _require_bound(self) -> None:
        """Raise a clear `RuntimeError` if this transform is unbound.

        Called as the first line of every artifact-consuming `apply`.
        """
        if not self.is_bound:
            raise RuntimeError(
                f"{type(self).__name__} was constructed with an ArtifactSource and is unbound; pass "
                f"it through a control that resolves and binds it during steer() (ActivationAdapter, "
                f"or CAST via behavior_transform)."
            )

    @property
    def covered_layer_ids(self) -> set[int] | None:
        """Layers this transform can act on; None = unknown (opts out of adapter validation).

        Artifact transforms report `set(directions.keys())` when bound and None when unbound;
        wrappers delegate to their inner transform.
        """
        return None

    def export(self, layer_id: int) -> "WireForm | None":
        """This configuration's wire form at `layer_id`, or None when the configuration is
        not expressible in the wire vocabulary.

        Exportability is a property of a configuration, not a class; an artifact whose shape
        has no wire semantics (e.g. a positional direction) returns None even though the
        class names a `wire_kind`. The default returns None.
        """
        return None

    def wire_plan(self) -> str | None:
        """The wire transform kind this configuration serializes to, or None when hook-only.

        Readable on the unbound form: a source-carrying transform consults its source's
        declared shape rather than resolving it. The default returns the class `wire_kind`.
        """
        return type(self).wire_kind

    def modifier_wire_kind(self, core_kind: str) -> str | None:
        """The wire modifier kind this wrapper contributes over a core transform kind, or
        None when the combination has no wire form.

        Meaningful only on wrapper transforms (`is_modifier` True). The default returns the
        class `wire_kind`.
        """
        return type(self).wire_kind

    def export_modifier(self, layer_id: int) -> "WireForm | None":
        """The wrapper's wire modifier form at `layer_id`, or None to contribute no modifier
        at that layer.

        Meaningful only on wrapper transforms (`is_modifier` True); kind-level
        inexpressibility is reported by `modifier_wire_kind`. The default returns None.
        """
        return None

artifact_meta property

Provenance metadata of the transform's steering artifact, or None.

Populated when the artifact was supplied or resolved as a SteeringVector carrying meta (fit fingerprints); consumers cross-check it against a serving engine's model identity. Wrappers delegate to their inner transform. The default returns None.

covered_layer_ids property

Layers this transform can act on; None = unknown (opts out of adapter validation).

Artifact transforms report set(directions.keys()) when bound and None when unbound; wrappers delegate to their inner transform.

is_bound property

True when the transform can be applied as-is (artifact concrete).

Default True: transforms that take no steering artifact are always bound.

is_modifier = False class-attribute

source property

The unresolved ArtifactSource this transform carries, or None when concrete.

wire_kind = None class-attribute

_require_bound()

Raise a clear RuntimeError if this transform is unbound.

Called as the first line of every artifact-consuming apply.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def _require_bound(self) -> None:
    """Raise a clear `RuntimeError` if this transform is unbound.

    Called as the first line of every artifact-consuming `apply`.
    """
    if not self.is_bound:
        raise RuntimeError(
            f"{type(self).__name__} was constructed with an ArtifactSource and is unbound; pass "
            f"it through a control that resolves and binds it during steer() (ActivationAdapter, "
            f"or CAST via behavior_transform)."
        )

apply(hidden_states, *, layer_id, token_mask, **kwargs) abstractmethod

Apply the transform and return modified hidden states.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Shape [B, T, H].	required
layer_id	int	Which layer this is being applied at.	required
token_mask	BoolTensor	Shape [B, T]. True at positions to modify.	required
**kwargs	Any	Transform-specific extra arguments.	{}

Returns:

Type	Description
Tensor	Modified hidden states, same shape as input.

Source code in steerability/algorithms/state_control/common/transforms/base.py

@abstractmethod
def apply(
    self,
    hidden_states: torch.Tensor,
    *,
    layer_id: int,
    token_mask: torch.BoolTensor,
    **kwargs: Any,
) -> torch.Tensor:
    """Apply the transform and return modified hidden states.

    Args:
        hidden_states: Shape [B, T, H].
        layer_id: Which layer this is being applied at.
        token_mask: Shape [B, T]. True at positions to modify.
        **kwargs: Transform-specific extra arguments.

    Returns:
        Modified hidden states, same shape as input.
    """
    ...

bind(ctx)

Return a fully-bound transform for this context.

Contract:

- MUST NOT mutate `self`; instances and sources are shared across adapters and
  `ControlSpec` grid points, whose params objects are reused per point.
- Returns `self` when already bound (idempotent).
- When source-carrying, returns a NEW instance of the same class constructed with
  `ctx.resolve(self._source)` and all hyperparameters copied, built by fresh
  construction rather than `copy.copy`.

Default: returns self.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def bind(self, ctx: "TransformContext") -> "BaseTransform":
    """Return a fully-bound transform for this context.

    Contract:

        - MUST NOT mutate `self`; instances and sources are shared across adapters and
          `ControlSpec` grid points, whose params objects are reused per point.
        - Returns `self` when already bound (idempotent).
        - When source-carrying, returns a NEW instance of the same class constructed with
          `ctx.resolve(self._source)` and all hyperparameters copied, built by fresh
          construction rather than `copy.copy`.

    Default: returns self.
    """
    return self

export(layer_id)

This configuration's wire form at layer_id, or None when the configuration is not expressible in the wire vocabulary.

Exportability is a property of a configuration, not a class; an artifact whose shape has no wire semantics (e.g. a positional direction) returns None even though the class names a wire_kind. The default returns None.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def export(self, layer_id: int) -> "WireForm | None":
    """This configuration's wire form at `layer_id`, or None when the configuration is
    not expressible in the wire vocabulary.

    Exportability is a property of a configuration, not a class; an artifact whose shape
    has no wire semantics (e.g. a positional direction) returns None even though the
    class names a `wire_kind`. The default returns None.
    """
    return None

export_modifier(layer_id)

The wrapper's wire modifier form at layer_id, or None to contribute no modifier at that layer.

Meaningful only on wrapper transforms (is_modifier True); kind-level inexpressibility is reported by modifier_wire_kind. The default returns None.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def export_modifier(self, layer_id: int) -> "WireForm | None":
    """The wrapper's wire modifier form at `layer_id`, or None to contribute no modifier
    at that layer.

    Meaningful only on wrapper transforms (`is_modifier` True); kind-level
    inexpressibility is reported by `modifier_wire_kind`. The default returns None.
    """
    return None

modifier_wire_kind(core_kind)

The wire modifier kind this wrapper contributes over a core transform kind, or None when the combination has no wire form.

Meaningful only on wrapper transforms (is_modifier True). The default returns the class wire_kind.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def modifier_wire_kind(self, core_kind: str) -> str | None:
    """The wire modifier kind this wrapper contributes over a core transform kind, or
    None when the combination has no wire form.

    Meaningful only on wrapper transforms (`is_modifier` True). The default returns the
    class `wire_kind`.
    """
    return type(self).wire_kind

wire_plan()

The wire transform kind this configuration serializes to, or None when hook-only.

Readable on the unbound form: a source-carrying transform consults its source's declared shape rather than resolving it. The default returns the class wire_kind.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def wire_plan(self) -> str | None:
    """The wire transform kind this configuration serializes to, or None when hook-only.

    Readable on the unbound form: a source-carrying transform consults its source's
    declared shape rather than resolving it. The default returns the class `wire_kind`.
    """
    return type(self).wire_kind

unwrap_modifiers(transform)

Split a possibly wrapper-chained transform into its core and its modifiers.

Parameters:

Name	Type	Description	Default
transform	BaseTransform	The transform, possibly wrapped in modifier transforms.	required

Returns:

Type	Description
BaseTransform	The core transform and the modifiers in application order, innermost wrapper first,
tuple[BaseTransform, ...]	matching the wire interpreter's composition order.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def unwrap_modifiers(transform: BaseTransform) -> tuple[BaseTransform, tuple[BaseTransform, ...]]:
    """Split a possibly wrapper-chained transform into its core and its modifiers.

    Args:
        transform: The transform, possibly wrapped in modifier transforms.

    Returns:
        The core transform and the modifiers in application order, innermost wrapper first,
        matching the wire interpreter's composition order.
    """
    wrappers: list[BaseTransform] = []
    current = transform
    while type(current).is_modifier:
        wrappers.append(current)
        current = current.inner
    return current, tuple(reversed(wrappers))

context

Context handed to transforms (and transform factories) at steer() time.

TransformContext dataclass

Fully-resolved layer/sizing context passed to a transform's bind or a factory.

A transform slot may be a BaseTransform instance (bound, or source-carrying and bound via bind(ctx)) or a factory Callable[[TransformContext], BaseTransform]. Either way it receives this context after the behavior layers are resolved, so it can turn a source into a concrete artifact (ctx.resolve(...)) and size any per-layer or per-head buffers to the model.

Transforms never touch the model directly: fit orchestration lives behind resolve, a closure the adapter builds over (model, tokenizer, device, dtype) that coerces the input, resolves the source (memoized master, cloned), and moves the clone onto the model's device/dtype.

Attributes:

Name	Type	Description
layer_ids	list[int]	The resolved behavior layer(s).
num_layers	int	Total number of transformer layers in the model.
hidden_size	int	Residual-stream width H.
num_heads	int | None	Number of attention heads per layer, or None if unavailable from the config.
head_dim	int | None	Dimension of each attention head, or None if unavailable from the config.
dtype	dtype	The model's parameter dtype.
device	device	The model's device.
resolve	Callable[[ArtifactSource | SteeringVector | Mapping[int, Tensor]], SteeringVector]	Turns a SteeringVector / directions mapping / ArtifactSource into a concrete SteeringVector on the model's device and dtype.

Source code in steerability/algorithms/state_control/common/transforms/context.py

@dataclass(frozen=True)
class TransformContext:
    """Fully-resolved layer/sizing context passed to a transform's `bind` or a factory.

    A transform slot may be a `BaseTransform` instance (bound, or source-carrying and bound via
    `bind(ctx)`) or a factory `Callable[[TransformContext], BaseTransform]`. Either way it receives
    this context after the behavior layers are resolved, so it can turn a source into a concrete
    artifact (`ctx.resolve(...)`) and size any per-layer or per-head buffers to the model.

    Transforms never touch the model directly: fit orchestration lives behind `resolve`, a closure
    the adapter builds over `(model, tokenizer, device, dtype)` that coerces the input, resolves the
    source (memoized master, cloned), and moves the clone onto the model's device/dtype.

    Attributes:
        layer_ids: The resolved behavior layer(s).
        num_layers: Total number of transformer layers in the model.
        hidden_size: Residual-stream width `H`.
        num_heads: Number of attention heads per layer, or None if unavailable from the config.
        head_dim: Dimension of each attention head, or None if unavailable from the config.
        dtype: The model's parameter dtype.
        device: The model's device.
        resolve: Turns a `SteeringVector` / directions mapping / `ArtifactSource` into a concrete
            `SteeringVector` on the model's device and dtype.
    """

    layer_ids: list[int]
    num_layers: int
    hidden_size: int
    num_heads: int | None
    head_dim: int | None
    dtype: torch.dtype
    device: torch.device
    resolve: Callable[[ArtifactSource | SteeringVector | Mapping[int, torch.Tensor]], SteeringVector]

device instance-attribute

dtype instance-attribute

head_dim instance-attribute

hidden_size instance-attribute

layer_ids instance-attribute

num_heads instance-attribute

num_layers instance-attribute

resolve instance-attribute

__init__(layer_ids, num_layers, hidden_size, num_heads, head_dim, dtype, device, resolve)

_build_context(model, tokenizer, layer_ids, layout=None, session=None)

Build the TransformContext for the given behavior layers.

With a live model, reads device/dtype/layer-count from the model and hidden_size from the text config; num_heads/head_dim come from the first behavior layer's head geometry read off the module tree (head_geometry), so a per-head transform reshapes with the geometry that layer actually uses. With no behavior layer given, or when no behavior layer carries an attention module (a residual-stream transform on the non-attention layers of a hybrid stack), the head geometry falls back to the text config. The resolve closure coerces any artifact to a source, resolves it against the model, and moves the result onto the model's device and dtype. With model=None, sizes come from layout (a structural core.execution.ModelFacts), the device is CPU, and the resolve closure serves concrete artifacts only, since fitting a source requires a live model.

Source code in steerability/algorithms/state_control/common/transforms/context.py

def _build_context(
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase | None,
    layer_ids: Sequence[int],
    layout=None,
    session=None,
) -> TransformContext:
    """Build the `TransformContext` for the given behavior layers.

    With a live model, reads device/dtype/layer-count from the model and `hidden_size` from the
    text config; `num_heads`/`head_dim` come from the first behavior layer's head geometry read
    off the module tree (`head_geometry`), so a per-head transform reshapes with the geometry that
    layer actually uses. With no behavior layer given, or when no behavior layer carries an
    attention module (a residual-stream transform on the non-attention layers of a hybrid stack),
    the head geometry falls back to the text config. The resolve closure coerces any artifact to a
    source, resolves it against the model,
    and moves the result onto the model's device and dtype. With `model=None`, sizes come from
    `layout` (a structural `core.execution.ModelFacts`), the device is CPU, and the resolve
    closure serves concrete artifacts only, since fitting a source requires a live model.
    """
    if model is not None:
        device = next(model.parameters()).device
        dtype = model.dtype
        module_layout = resolve_model_layout(model)
        num_layers = module_layout.num_layers

        text_cfg = text_config(model)
        hidden_size = text_cfg.hidden_size
        attention_layer = next((lid for lid in layer_ids if module_layout.has_attention(lid)), None)
        if attention_layer is not None:
            geometry = head_geometry(model, module_layout, attention_layer)
            num_heads = geometry.num_heads
            head_dim = geometry.head_dim
        else:
            num_heads = getattr(text_cfg, "num_attention_heads", None)
            head_dim = getattr(text_cfg, "head_dim", None)
            if head_dim is None and num_heads:
                head_dim = hidden_size // num_heads
    else:
        if layout is None:
            raise ValueError("Building a TransformContext requires a live model or a structural layout.")
        device = torch.device("cpu")
        dtype = getattr(torch, layout.dtype, None)
        if not isinstance(dtype, torch.dtype):
            raise ValueError(f"Layout dtype {layout.dtype!r} does not name a torch dtype.")
        num_layers = layout.num_layers
        hidden_size = layout.hidden_size
        num_heads = layout.num_attention_heads
        head_dim = layout.head_dim

    def resolve(artifact) -> SteeringVector:
        source = _as_artifact_source(artifact)
        try:
            resolved = source.resolve(model, tokenizer, session=session)
        except TypeError:
            resolved = source.resolve(model, tokenizer)  # sources without capture support
        return resolved.to(device, dtype)

    return TransformContext(
        layer_ids=list(layer_ids),
        num_layers=num_layers,
        hidden_size=hidden_size,
        num_heads=num_heads,
        head_dim=head_dim,
        dtype=dtype,
        device=device,
        resolve=resolve,
    )

resolve_transform_slot(slot, model, tokenizer, layer_ids, layout=None, require_coverage=True, session=None)

Turn a transform slot into a bound, coverage-checked BaseTransform for the given model.

Shared by every control that exposes a transform slot (ActivationAdapter.transform, CAST's behavior_transform), so the resolution contract lives in one place. It runs three stages:

1. Context construction: introspect the model (device, dtype, layer count, hidden size, head count/dim) and build a TransformContext whose resolve closure fits/resolves a source on the steered model and moves the result onto its device/dtype (_build_context).
2. Slot resolution: a BaseTransform instance is bound if unbound (else used as-is); a factory is invoked with the context and its result is bound if unbound. Either branch is validated to yield a bound BaseTransform, so a factory returning a source-carrying transform is bound here rather than failing later in apply.
3. Coverage check: when the resolved transform reports covered_layer_ids, every layer in layer_ids must be covered; a transform reporting None opts out of the check.

Parameters:

Name	Type	Description	Default
slot	BaseTransform | Callable[[TransformContext], BaseTransform]	A BaseTransform (bound or source-carrying) or a factory taking the context.	required
model	PreTrainedModel | None	The steered model to introspect and resolve sources against, or None to build the context from layout (concrete artifacts only; fitting a source requires a model).	required
tokenizer	PreTrainedTokenizerBase | None	Tokenizer used when a source fits from data; may be None for concrete artifacts.	required
layer_ids	Sequence[int]	The resolved behavior layers the transform must cover.	required
layout	ModelFacts | None	Structural core.execution.ModelFacts consulted when model is None.	None
require_coverage	bool	When False, skip the coverage check; uncovered layers are hooked and pass through unchanged at apply time.	True
session	SteeringSession | None	Optional SteeringSession forwarded to sources whose fit runs through session capture.	None

Returns:

Type	Description
BaseTransform	A bound BaseTransform ready for apply.

Raises:

Type	Description
TypeError	If a factory returns a non-BaseTransform, or a bind returns an unbound transform.
ValueError	If the transform covers only some of layer_ids.

Source code in steerability/algorithms/state_control/common/transforms/context.py

def resolve_transform_slot(
    slot: BaseTransform | Callable[[TransformContext], BaseTransform],
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase | None,
    layer_ids: Sequence[int],
    layout: ModelFacts | None = None,
    require_coverage: bool = True,
    session: SteeringSession | None = None,
) -> BaseTransform:
    """Turn a transform slot into a bound, coverage-checked `BaseTransform` for the given model.

    Shared by every control that exposes a transform slot (`ActivationAdapter.transform`, `CAST`'s
    `behavior_transform`), so the resolution contract lives in one place. It runs three stages:

    1. **Context construction**: introspect the model (device, dtype, layer count, hidden size,
       head count/dim) and build a `TransformContext` whose `resolve` closure fits/resolves a source
       on the steered model and moves the result onto its device/dtype (`_build_context`).
    2. **Slot resolution**: a `BaseTransform` instance is bound if unbound (else used as-is); a
       factory is invoked with the context and its result is bound if unbound. Either branch is
       validated to yield a bound `BaseTransform`, so a factory returning a source-carrying transform
       is bound here rather than failing later in `apply`.
    3. **Coverage check**: when the resolved transform reports `covered_layer_ids`, every layer in
       `layer_ids` must be covered; a transform reporting `None` opts out of the check.

    Args:
        slot: A `BaseTransform` (bound or source-carrying) or a factory taking the context.
        model: The steered model to introspect and resolve sources against, or None to build the
            context from `layout` (concrete artifacts only; fitting a source requires a model).
        tokenizer: Tokenizer used when a source fits from data; may be None for concrete artifacts.
        layer_ids: The resolved behavior layers the transform must cover.
        layout: Structural `core.execution.ModelFacts` consulted when `model` is None.
        require_coverage: When False, skip the coverage check; uncovered layers are hooked and
            pass through unchanged at apply time.
        session: Optional `SteeringSession` forwarded to sources whose fit runs through
            session capture.

    Returns:
        A bound `BaseTransform` ready for `apply`.

    Raises:
        TypeError: If a factory returns a non-`BaseTransform`, or a `bind` returns an unbound
            transform.
        ValueError: If the transform covers only some of `layer_ids`.
    """
    ctx = _build_context(model, tokenizer, layer_ids, layout=layout, session=session)

    if isinstance(slot, BaseTransform):
        built = slot if slot.is_bound else slot.bind(ctx)
    else:  # factory
        built = slot(ctx)
        if not isinstance(built, BaseTransform):
            raise TypeError(f"transform factory must return a BaseTransform; got {type(built).__name__}.")
        if not built.is_bound:
            built = built.bind(ctx)

    if not isinstance(built, BaseTransform) or not built.is_bound:
        raise TypeError(
            f"{type(slot).__name__}.bind() must return a bound BaseTransform; got "
            f"{type(built).__name__} (is_bound={getattr(built, 'is_bound', None)!r})."
        )

    coverage = built.covered_layer_ids
    if require_coverage and coverage is not None:
        missing = [lid for lid in layer_ids if lid not in coverage]
        if missing:
            raise ValueError(
                f"Transform has no direction for layer(s) {missing} (covers {sorted(coverage)})."
            )

    return built

head_additive

Head-level additive transform for activation steering.

HeadAdditiveTransform

Bases: BaseTransform

Adds scaled direction vectors to specific head slices.

For each selected (layer, head) pair, it adds a direction vector to the slice [h * head_dim : (h+1) * head_dim].

For ITI, this operates in pre-o_proj space: the input to the output projection where each head_dim-sized slice corresponds to an individual attention head's output. The directions must be computed in the same space.

Expects a SteeringVector whose directions are shaped [num_heads, head_dim] per layer, with num_heads and head_dim metadata set. Only head indices present in active_heads are applied; other heads are left untouched.

A bare directions mapping cannot carry the required num_heads/head_dim metadata and is rejected; supply a SteeringVector (or an ArtifactSource resolving to one).

Parameters:

Name	Type	Description	Default
artifact	SteeringVector | Mapping[int, Tensor] | ArtifactSource	A SteeringVector with per-head directions and num_heads/head_dim metadata, or an ArtifactSource (unbound until bind(ctx)). Required.	required
active_heads	dict[int, set[int]]	Mapping from layer_id to set of head indices to intervene on.	required
strength	float	Global scaling factor (alpha in ITI).	1.0

Source code in steerability/algorithms/state_control/common/transforms/head_additive.py

class HeadAdditiveTransform(BaseTransform):
    """Adds scaled direction vectors to specific head slices.

    For each selected (layer, head) pair, it adds a direction vector to the slice
    [h * head_dim : (h+1) * head_dim].

    For ITI, this operates in pre-o_proj space: the input to the output projection
    where each head_dim-sized slice corresponds to an individual attention head's
    output. The directions must be computed in the same space.

    Expects a SteeringVector whose directions are shaped [num_heads, head_dim]
    per layer, with num_heads and head_dim metadata set. Only head indices
    present in ``active_heads`` are applied; other heads are left untouched.

    A bare directions mapping cannot carry the required `num_heads`/`head_dim` metadata and is
    rejected; supply a `SteeringVector` (or an `ArtifactSource` resolving to one).

    Args:
        artifact: A `SteeringVector` with per-head directions and `num_heads`/`head_dim` metadata,
            or an `ArtifactSource` (unbound until `bind(ctx)`). Required.
        active_heads: Mapping from layer_id to set of head indices to intervene on.
        strength: Global scaling factor (alpha in ITI).
    """

    wire_kind: ClassVar[str | None] = "head_additive"

    def __init__(
        self,
        artifact: SteeringVector | Mapping[int, torch.Tensor] | ArtifactSource,
        active_heads: dict[int, set[int]],
        strength: float = 1.0,
    ):
        self.active_heads = active_heads
        self.strength = strength
        self._source: ArtifactSource | None = None
        self.steering_vector: SteeringVector | None = None
        self.num_heads: int | None = None
        self.head_dim: int | None = None

        if isinstance(artifact, ArtifactSource):
            self._source = artifact
        elif isinstance(artifact, SteeringVector):
            self.steering_vector = artifact
            self._validate_artifact()
        elif isinstance(artifact, Mapping):
            raise ValueError(
                "HeadAdditiveTransform requires num_heads and head_dim metadata on the SteeringVector; "
                "a bare directions mapping cannot carry it."
            )
        else:
            raise TypeError(
                f"HeadAdditiveTransform artifact must be a SteeringVector or an ArtifactSource; got "
                f"{type(artifact).__name__} (did you mean strength=?)."
            )

    def _validate_artifact(self) -> None:
        """Validate the artifact carries num_heads/head_dim metadata."""
        if self.steering_vector.num_heads is None or self.steering_vector.head_dim is None:
            raise ValueError("HeadAdditiveTransform requires num_heads and head_dim metadata on the SteeringVector.")
        self.num_heads = self.steering_vector.num_heads
        self.head_dim = self.steering_vector.head_dim

    @property
    def is_bound(self) -> bool:
        return self.steering_vector is not None

    @property
    def artifact_meta(self) -> dict | None:
        if self.steering_vector is not None and self.steering_vector.meta:
            return dict(self.steering_vector.meta)
        return None

    def bind(self, ctx: "TransformContext") -> "HeadAdditiveTransform":
        if self.is_bound:
            return self
        return HeadAdditiveTransform(ctx.resolve(self._source), active_heads=self.active_heads, strength=self.strength)

    @property
    def covered_layer_ids(self) -> set[int] | None:
        """Layers the transform can act on: layers with directions and active heads."""
        if self.steering_vector is None:
            return None
        return {
            layer_id
            for layer_id in self.steering_vector.directions
            if self.active_heads.get(layer_id)
        }

    def to_config(self) -> tuple[dict, object | None, "BaseTransform | None"]:
        """The `(params, artifact, inner)` serialized form under the `head_additive` kind.

        `active_heads` serializes with string layer keys and sorted head lists.
        """
        params = {
            "strength": float(self.strength),
            "active_heads": {str(layer): sorted(int(h) for h in heads)
                             for layer, heads in self.active_heads.items()},
        }
        artifact = self.steering_vector if self.steering_vector is not None else self._source
        return params, artifact, None

    @classmethod
    def from_config(cls, params: dict, *, artifact=None, inner=None) -> "HeadAdditiveTransform":
        """Rebuild a `head_additive` transform from its serialized form."""
        active_heads = {int(layer): set(heads) for layer, heads in params.get("active_heads", {}).items()}
        return cls(artifact, active_heads=active_heads, strength=params.get("strength", 1.0))

    def export(self, layer_id: int) -> "WireForm | None":
        """The `head_additive` wire form for `layer_id`.

        The wire vector is `[num_heads, head_dim]` with zeros at heads outside `active_heads`,
        so the broadcast wire addition reproduces the selective per-head addition exactly.
        """
        from ..specs import WireForm

        if self.steering_vector is None:
            return None
        heads = self.active_heads.get(layer_id)
        dirs = self.steering_vector.directions.get(layer_id)
        if not heads or dirs is None:
            return None
        vector = torch.zeros(self.num_heads, self.head_dim, dtype=dirs.dtype)
        for head_id in heads:
            vector[head_id] = dirs[head_id]
        return WireForm(
            kind="head_additive",
            params={"strength": float(self.strength)},
            tensors={"vector": vector},
        )


    def apply(
        self,
        hidden_states: torch.Tensor,
        *,
        layer_id: int,
        token_mask: torch.BoolTensor,
        **kwargs,
    ) -> torch.Tensor:
        """Apply head-level additive steering.

        Args:
            hidden_states: Shape [B, T, H] where H = num_heads * head_dim.
            layer_id: Which layer this is being applied at.
            token_mask: Shape [B, T]. True at positions to modify.

        Returns:
            Modified hidden states, same shape as input.
        """
        self._require_bound()
        heads = self.active_heads.get(layer_id)
        if not heads:
            return hidden_states

        dirs = self.steering_vector.directions.get(layer_id)
        if dirs is None:
            return hidden_states

        hidden_states = hidden_states.clone()

        for head_id in heads:
            start = head_id * self.head_dim
            end = start + self.head_dim
            direction = dirs[head_id]  # [head_dim]
            v = (self.strength * direction).to(dtype=hidden_states.dtype, device=hidden_states.device)
            # scale by token mask so unmasked positions are untouched
            delta = token_mask.unsqueeze(-1).to(hidden_states.dtype) * v.view(1, 1, -1)
            hidden_states[:, :, start:end] = hidden_states[:, :, start:end] + delta

        return hidden_states

_source = None instance-attribute

active_heads = active_heads instance-attribute

artifact_meta property

Provenance metadata of the transform's steering artifact, or None.

Populated when the artifact was supplied or resolved as a SteeringVector carrying meta (fit fingerprints); consumers cross-check it against a serving engine's model identity. Wrappers delegate to their inner transform. The default returns None.

covered_layer_ids property

Layers the transform can act on: layers with directions and active heads.

head_dim = None instance-attribute

is_bound property

True when the transform can be applied as-is (artifact concrete).

Default True: transforms that take no steering artifact are always bound.

is_modifier = False class-attribute

num_heads = None instance-attribute

source property

The unresolved ArtifactSource this transform carries, or None when concrete.

steering_vector = None instance-attribute

strength = strength instance-attribute

wire_kind = 'head_additive' class-attribute

__init__(artifact, active_heads, strength=1.0)

Source code in steerability/algorithms/state_control/common/transforms/head_additive.py

def __init__(
    self,
    artifact: SteeringVector | Mapping[int, torch.Tensor] | ArtifactSource,
    active_heads: dict[int, set[int]],
    strength: float = 1.0,
):
    self.active_heads = active_heads
    self.strength = strength
    self._source: ArtifactSource | None = None
    self.steering_vector: SteeringVector | None = None
    self.num_heads: int | None = None
    self.head_dim: int | None = None

    if isinstance(artifact, ArtifactSource):
        self._source = artifact
    elif isinstance(artifact, SteeringVector):
        self.steering_vector = artifact
        self._validate_artifact()
    elif isinstance(artifact, Mapping):
        raise ValueError(
            "HeadAdditiveTransform requires num_heads and head_dim metadata on the SteeringVector; "
            "a bare directions mapping cannot carry it."
        )
    else:
        raise TypeError(
            f"HeadAdditiveTransform artifact must be a SteeringVector or an ArtifactSource; got "
            f"{type(artifact).__name__} (did you mean strength=?)."
        )

_require_bound()

Raise a clear RuntimeError if this transform is unbound.

Called as the first line of every artifact-consuming apply.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def _require_bound(self) -> None:
    """Raise a clear `RuntimeError` if this transform is unbound.

    Called as the first line of every artifact-consuming `apply`.
    """
    if not self.is_bound:
        raise RuntimeError(
            f"{type(self).__name__} was constructed with an ArtifactSource and is unbound; pass "
            f"it through a control that resolves and binds it during steer() (ActivationAdapter, "
            f"or CAST via behavior_transform)."
        )

_validate_artifact()

Validate the artifact carries num_heads/head_dim metadata.

Source code in steerability/algorithms/state_control/common/transforms/head_additive.py

def _validate_artifact(self) -> None:
    """Validate the artifact carries num_heads/head_dim metadata."""
    if self.steering_vector.num_heads is None or self.steering_vector.head_dim is None:
        raise ValueError("HeadAdditiveTransform requires num_heads and head_dim metadata on the SteeringVector.")
    self.num_heads = self.steering_vector.num_heads
    self.head_dim = self.steering_vector.head_dim

apply(hidden_states, *, layer_id, token_mask, **kwargs)

Apply head-level additive steering.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Shape [B, T, H] where H = num_heads * head_dim.	required
layer_id	int	Which layer this is being applied at.	required
token_mask	BoolTensor	Shape [B, T]. True at positions to modify.	required

Returns:

Type	Description
Tensor	Modified hidden states, same shape as input.

Source code in steerability/algorithms/state_control/common/transforms/head_additive.py

def apply(
    self,
    hidden_states: torch.Tensor,
    *,
    layer_id: int,
    token_mask: torch.BoolTensor,
    **kwargs,
) -> torch.Tensor:
    """Apply head-level additive steering.

    Args:
        hidden_states: Shape [B, T, H] where H = num_heads * head_dim.
        layer_id: Which layer this is being applied at.
        token_mask: Shape [B, T]. True at positions to modify.

    Returns:
        Modified hidden states, same shape as input.
    """
    self._require_bound()
    heads = self.active_heads.get(layer_id)
    if not heads:
        return hidden_states

    dirs = self.steering_vector.directions.get(layer_id)
    if dirs is None:
        return hidden_states

    hidden_states = hidden_states.clone()

    for head_id in heads:
        start = head_id * self.head_dim
        end = start + self.head_dim
        direction = dirs[head_id]  # [head_dim]
        v = (self.strength * direction).to(dtype=hidden_states.dtype, device=hidden_states.device)
        # scale by token mask so unmasked positions are untouched
        delta = token_mask.unsqueeze(-1).to(hidden_states.dtype) * v.view(1, 1, -1)
        hidden_states[:, :, start:end] = hidden_states[:, :, start:end] + delta

    return hidden_states

bind(ctx)

Return a fully-bound transform for this context.

Contract:

- MUST NOT mutate `self`; instances and sources are shared across adapters and
  `ControlSpec` grid points, whose params objects are reused per point.
- Returns `self` when already bound (idempotent).
- When source-carrying, returns a NEW instance of the same class constructed with
  `ctx.resolve(self._source)` and all hyperparameters copied, built by fresh
  construction rather than `copy.copy`.

Default: returns self.

Source code in steerability/algorithms/state_control/common/transforms/head_additive.py

def bind(self, ctx: "TransformContext") -> "HeadAdditiveTransform":
    if self.is_bound:
        return self
    return HeadAdditiveTransform(ctx.resolve(self._source), active_heads=self.active_heads, strength=self.strength)

export(layer_id)

The head_additive wire form for layer_id.

The wire vector is [num_heads, head_dim] with zeros at heads outside active_heads, so the broadcast wire addition reproduces the selective per-head addition exactly.

Source code in steerability/algorithms/state_control/common/transforms/head_additive.py

def export(self, layer_id: int) -> "WireForm | None":
    """The `head_additive` wire form for `layer_id`.

    The wire vector is `[num_heads, head_dim]` with zeros at heads outside `active_heads`,
    so the broadcast wire addition reproduces the selective per-head addition exactly.
    """
    from ..specs import WireForm

    if self.steering_vector is None:
        return None
    heads = self.active_heads.get(layer_id)
    dirs = self.steering_vector.directions.get(layer_id)
    if not heads or dirs is None:
        return None
    vector = torch.zeros(self.num_heads, self.head_dim, dtype=dirs.dtype)
    for head_id in heads:
        vector[head_id] = dirs[head_id]
    return WireForm(
        kind="head_additive",
        params={"strength": float(self.strength)},
        tensors={"vector": vector},
    )

export_modifier(layer_id)

The wrapper's wire modifier form at layer_id, or None to contribute no modifier at that layer.

Meaningful only on wrapper transforms (is_modifier True); kind-level inexpressibility is reported by modifier_wire_kind. The default returns None.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def export_modifier(self, layer_id: int) -> "WireForm | None":
    """The wrapper's wire modifier form at `layer_id`, or None to contribute no modifier
    at that layer.

    Meaningful only on wrapper transforms (`is_modifier` True); kind-level
    inexpressibility is reported by `modifier_wire_kind`. The default returns None.
    """
    return None

from_config(params, *, artifact=None, inner=None) classmethod

Rebuild a head_additive transform from its serialized form.

Source code in steerability/algorithms/state_control/common/transforms/head_additive.py

@classmethod
def from_config(cls, params: dict, *, artifact=None, inner=None) -> "HeadAdditiveTransform":
    """Rebuild a `head_additive` transform from its serialized form."""
    active_heads = {int(layer): set(heads) for layer, heads in params.get("active_heads", {}).items()}
    return cls(artifact, active_heads=active_heads, strength=params.get("strength", 1.0))

modifier_wire_kind(core_kind)

The wire modifier kind this wrapper contributes over a core transform kind, or None when the combination has no wire form.

Meaningful only on wrapper transforms (is_modifier True). The default returns the class wire_kind.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def modifier_wire_kind(self, core_kind: str) -> str | None:
    """The wire modifier kind this wrapper contributes over a core transform kind, or
    None when the combination has no wire form.

    Meaningful only on wrapper transforms (`is_modifier` True). The default returns the
    class `wire_kind`.
    """
    return type(self).wire_kind

to_config()

The (params, artifact, inner) serialized form under the head_additive kind.

active_heads serializes with string layer keys and sorted head lists.

Source code in steerability/algorithms/state_control/common/transforms/head_additive.py

def to_config(self) -> tuple[dict, object | None, "BaseTransform | None"]:
    """The `(params, artifact, inner)` serialized form under the `head_additive` kind.

    `active_heads` serializes with string layer keys and sorted head lists.
    """
    params = {
        "strength": float(self.strength),
        "active_heads": {str(layer): sorted(int(h) for h in heads)
                         for layer, heads in self.active_heads.items()},
    }
    artifact = self.steering_vector if self.steering_vector is not None else self._source
    return params, artifact, None

wire_plan()

The wire transform kind this configuration serializes to, or None when hook-only.

Readable on the unbound form: a source-carrying transform consults its source's declared shape rather than resolving it. The default returns the class wire_kind.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def wire_plan(self) -> str | None:
    """The wire transform kind this configuration serializes to, or None when hook-only.

    Readable on the unbound form: a source-carrying transform consults its source's
    declared shape rather than resolving it. The default returns the class `wire_kind`.
    """
    return type(self).wire_kind

norm_preserving

Wrapper that rescales hidden states to preserve original norms.

NormPreservingTransform

Bases: BaseTransform

Wraps an inner transform and rescales to maintain original norms.

After applying the inner transform, if the norm increased at any position, rescale those positions back to original norm. This prevents distribution shift from large steering vectors.

Binding and coverage delegate to the inner transform: the wrapper is bound iff the inner is, bind returns a new wrapper around the bound inner, and coverage is the inner's coverage.

Parameters:

Name	Type	Description	Default
inner	BaseTransform	The transform to wrap.	required

Source code in steerability/algorithms/state_control/common/transforms/norm_preserving.py

class NormPreservingTransform(BaseTransform):
    """Wraps an inner transform and rescales to maintain original norms.

    After applying the inner transform, if the norm increased at any position,
    rescale those positions back to original norm. This prevents distribution
    shift from large steering vectors.

    Binding and coverage delegate to the inner transform: the wrapper is bound iff the inner is,
    `bind` returns a new wrapper around the bound inner, and coverage is the inner's coverage.

    Args:
        inner: The transform to wrap.
    """

    wire_kind: ClassVar[str | None] = "norm_preserving"
    is_modifier: ClassVar[bool] = True

    def __init__(self, inner: BaseTransform):
        self._inner = inner

    @property
    def inner(self) -> BaseTransform:
        """The wrapped transform."""
        return self._inner

    @property
    def is_bound(self) -> bool:
        return self._inner.is_bound

    @property
    def artifact_meta(self) -> dict | None:
        return self._inner.artifact_meta

    def bind(self, ctx: "TransformContext") -> "NormPreservingTransform":
        if self.is_bound:
            return self
        return NormPreservingTransform(self._inner.bind(ctx))

    @property
    def covered_layer_ids(self) -> set[int] | None:
        return self._inner.covered_layer_ids

    def to_config(self) -> tuple[dict, object | None, BaseTransform | None]:
        """The `(params, artifact, inner)` serialized form under the `norm_preserving` kind."""
        return {}, None, self._inner

    @classmethod
    def from_config(cls, params: dict, *, artifact=None, inner=None) -> "NormPreservingTransform":
        """Rebuild a `norm_preserving` wrapper around its decoded inner transform."""
        return cls(inner)

    def modifier_wire_kind(self, core_kind: str) -> str | None:
        """`"norm_preserving"`, or None over a per-head core.

        The wire modifier rescales over the last tensor dimension, which matches the hook
        semantics on the residual stream only; a wrapped `head_additive` (per-head stream) is
        hook-only.
        """
        if core_kind == "head_additive":
            return None
        return "norm_preserving"

    def export_modifier(self, layer_id: int) -> "WireForm | None":
        """The `norm_preserving` wire modifier form (no params, no tensors)."""
        from ..specs import WireForm

        return WireForm(kind="norm_preserving")


    def apply(
        self,
        hidden_states: torch.Tensor,
        *,
        layer_id: int,
        token_mask: torch.BoolTensor,
        **kwargs,
    ) -> torch.Tensor:
        """Apply inner transform then rescale to preserve norms.

        Args:
            hidden_states: Shape [B, T, H].
            layer_id: Which layer this is being applied at.
            token_mask: Shape [B, T]. True at positions to modify.

        Returns:
            Modified hidden states with preserved norms.

        Raises:
            ValueError: If NaN or Inf detected after transform.
        """
        self._require_bound()
        original_norm = hidden_states.norm(dim=-1, keepdim=True)
        modified = self._inner.apply(
            hidden_states, layer_id=layer_id, token_mask=token_mask, **kwargs
        )

        if torch.isnan(modified).any() or torch.isinf(modified).any():
            raise ValueError("NaN or Inf detected after transform application.")

        new_norm = modified.norm(dim=-1, keepdim=True)
        # only rescale where norm increased
        needs_rescale = new_norm > original_norm
        if needs_rescale.any():
            scale = torch.where(needs_rescale, original_norm / (new_norm + 1e-8), torch.ones_like(new_norm))
            modified = modified * scale

        return modified

_inner = inner instance-attribute

artifact_meta property

Provenance metadata of the transform's steering artifact, or None.

Populated when the artifact was supplied or resolved as a SteeringVector carrying meta (fit fingerprints); consumers cross-check it against a serving engine's model identity. Wrappers delegate to their inner transform. The default returns None.

covered_layer_ids property

Layers this transform can act on; None = unknown (opts out of adapter validation).

Artifact transforms report set(directions.keys()) when bound and None when unbound; wrappers delegate to their inner transform.

inner property

The wrapped transform.

is_bound property

True when the transform can be applied as-is (artifact concrete).

Default True: transforms that take no steering artifact are always bound.

is_modifier = True class-attribute

source property

The unresolved ArtifactSource this transform carries, or None when concrete.

wire_kind = 'norm_preserving' class-attribute

__init__(inner)

Source code in steerability/algorithms/state_control/common/transforms/norm_preserving.py

def __init__(self, inner: BaseTransform):
    self._inner = inner

_require_bound()

Raise a clear RuntimeError if this transform is unbound.

Called as the first line of every artifact-consuming apply.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def _require_bound(self) -> None:
    """Raise a clear `RuntimeError` if this transform is unbound.

    Called as the first line of every artifact-consuming `apply`.
    """
    if not self.is_bound:
        raise RuntimeError(
            f"{type(self).__name__} was constructed with an ArtifactSource and is unbound; pass "
            f"it through a control that resolves and binds it during steer() (ActivationAdapter, "
            f"or CAST via behavior_transform)."
        )

apply(hidden_states, *, layer_id, token_mask, **kwargs)

Apply inner transform then rescale to preserve norms.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Shape [B, T, H].	required
layer_id	int	Which layer this is being applied at.	required
token_mask	BoolTensor	Shape [B, T]. True at positions to modify.	required

Returns:

Type	Description
Tensor	Modified hidden states with preserved norms.

Raises:

Type	Description
ValueError	If NaN or Inf detected after transform.

Source code in steerability/algorithms/state_control/common/transforms/norm_preserving.py

def apply(
    self,
    hidden_states: torch.Tensor,
    *,
    layer_id: int,
    token_mask: torch.BoolTensor,
    **kwargs,
) -> torch.Tensor:
    """Apply inner transform then rescale to preserve norms.

    Args:
        hidden_states: Shape [B, T, H].
        layer_id: Which layer this is being applied at.
        token_mask: Shape [B, T]. True at positions to modify.

    Returns:
        Modified hidden states with preserved norms.

    Raises:
        ValueError: If NaN or Inf detected after transform.
    """
    self._require_bound()
    original_norm = hidden_states.norm(dim=-1, keepdim=True)
    modified = self._inner.apply(
        hidden_states, layer_id=layer_id, token_mask=token_mask, **kwargs
    )

    if torch.isnan(modified).any() or torch.isinf(modified).any():
        raise ValueError("NaN or Inf detected after transform application.")

    new_norm = modified.norm(dim=-1, keepdim=True)
    # only rescale where norm increased
    needs_rescale = new_norm > original_norm
    if needs_rescale.any():
        scale = torch.where(needs_rescale, original_norm / (new_norm + 1e-8), torch.ones_like(new_norm))
        modified = modified * scale

    return modified

bind(ctx)

Return a fully-bound transform for this context.

Contract:

- MUST NOT mutate `self`; instances and sources are shared across adapters and
  `ControlSpec` grid points, whose params objects are reused per point.
- Returns `self` when already bound (idempotent).
- When source-carrying, returns a NEW instance of the same class constructed with
  `ctx.resolve(self._source)` and all hyperparameters copied, built by fresh
  construction rather than `copy.copy`.

Default: returns self.

Source code in steerability/algorithms/state_control/common/transforms/norm_preserving.py

def bind(self, ctx: "TransformContext") -> "NormPreservingTransform":
    if self.is_bound:
        return self
    return NormPreservingTransform(self._inner.bind(ctx))

export(layer_id)

This configuration's wire form at layer_id, or None when the configuration is not expressible in the wire vocabulary.

Exportability is a property of a configuration, not a class; an artifact whose shape has no wire semantics (e.g. a positional direction) returns None even though the class names a wire_kind. The default returns None.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def export(self, layer_id: int) -> "WireForm | None":
    """This configuration's wire form at `layer_id`, or None when the configuration is
    not expressible in the wire vocabulary.

    Exportability is a property of a configuration, not a class; an artifact whose shape
    has no wire semantics (e.g. a positional direction) returns None even though the
    class names a `wire_kind`. The default returns None.
    """
    return None

export_modifier(layer_id)

The norm_preserving wire modifier form (no params, no tensors).

Source code in steerability/algorithms/state_control/common/transforms/norm_preserving.py

def export_modifier(self, layer_id: int) -> "WireForm | None":
    """The `norm_preserving` wire modifier form (no params, no tensors)."""
    from ..specs import WireForm

    return WireForm(kind="norm_preserving")

from_config(params, *, artifact=None, inner=None) classmethod

Rebuild a norm_preserving wrapper around its decoded inner transform.

Source code in steerability/algorithms/state_control/common/transforms/norm_preserving.py

@classmethod
def from_config(cls, params: dict, *, artifact=None, inner=None) -> "NormPreservingTransform":
    """Rebuild a `norm_preserving` wrapper around its decoded inner transform."""
    return cls(inner)

modifier_wire_kind(core_kind)

"norm_preserving", or None over a per-head core.

The wire modifier rescales over the last tensor dimension, which matches the hook semantics on the residual stream only; a wrapped head_additive (per-head stream) is hook-only.

Source code in steerability/algorithms/state_control/common/transforms/norm_preserving.py

def modifier_wire_kind(self, core_kind: str) -> str | None:
    """`"norm_preserving"`, or None over a per-head core.

    The wire modifier rescales over the last tensor dimension, which matches the hook
    semantics on the residual stream only; a wrapped `head_additive` (per-head stream) is
    hook-only.
    """
    if core_kind == "head_additive":
        return None
    return "norm_preserving"

to_config()

The (params, artifact, inner) serialized form under the norm_preserving kind.

Source code in steerability/algorithms/state_control/common/transforms/norm_preserving.py

def to_config(self) -> tuple[dict, object | None, BaseTransform | None]:
    """The `(params, artifact, inner)` serialized form under the `norm_preserving` kind."""
    return {}, None, self._inner

wire_plan()

The wire transform kind this configuration serializes to, or None when hook-only.

Readable on the unbound form: a source-carrying transform consults its source's declared shape rather than resolving it. The default returns the class wire_kind.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def wire_plan(self) -> str | None:
    """The wire transform kind this configuration serializes to, or None when hook-only.

    Readable on the unbound form: a source-carrying transform consults its source's
    declared shape rather than resolving it. The default returns the class `wire_kind`.
    """
    return type(self).wire_kind

projection

Projection transform: projects learned directions out of the residual stream.

ProjectionTransform

Bases: BaseTransform

Removes one or more learned directions from hidden states by projection.

For an orthonormal set of directions {d_1..d_k} at a layer (rows of a [K, H] tensor):

h' = h - alpha * sum_i (h . d_i) d_i          (applied at masked positions)

K=1 is single-direction ablation (the abliteration / directional-ablation technique of Arditi et al.); K>1 ablates the whole subspace. alpha in [0, 1] scales the removal: 1.0 fully removes the component (h'.d_i == 0), values < 1.0 give graded partial suppression.

This is a projection, not a translation or a rotation:

• It is idempotent at alpha=1 (P^2 = P).
• It is norm-reducing, since it drops a component.

The stored rows need not be orthonormal; they are orthonormalized via Gram-Schmidt and cached per (layer_id, device, dtype) so K>1 removal is basis-correct and order-independent.

Parameters:

Name	Type	Description	Default
artifact	SteeringVector | Mapping[int, Tensor] | ArtifactSource	The steering artifact, given as a SteeringVector, a per-layer directions mapping (Mapping[int, Tensor], each [K, H] or [H] treated as K=1), or an ArtifactSource (unbound until bind(ctx)). Required.	required
alpha	float	Ablation strength in [0, 1]. 1.0 = full removal (default); < 1.0 = partial.	1.0

Reference:

• "Refusal in Language Models Is Mediated by a Single Direction" Andy Arditi, Oscar Obeso, Aaquib Syed, Daniel Paleka, Nina Panickssery, Wes Gurnee, Neel Nanda https://arxiv.org/abs/2406.11717

Source code in steerability/algorithms/state_control/common/transforms/projection.py

class ProjectionTransform(BaseTransform):
    """Removes one or more learned directions from hidden states by projection.

    For an orthonormal set of directions `{d_1..d_k}` at a layer (rows of a `[K, H]` tensor):

        h' = h - alpha * sum_i (h . d_i) d_i          (applied at masked positions)

    `K=1` is single-direction ablation (the abliteration / directional-ablation technique of
    Arditi et al.); `K>1` ablates the whole subspace. `alpha` in `[0, 1]` scales the removal:
    `1.0` fully removes the component (`h'.d_i == 0`), values `< 1.0` give graded partial
    suppression.

    This is a projection, not a translation or a rotation:

    - It is idempotent at `alpha=1` (`P^2 = P`).
    - It is norm-reducing, since it drops a component.

    The stored rows need not be orthonormal; they are orthonormalized via Gram-Schmidt and cached
    per `(layer_id, device, dtype)` so `K>1` removal is basis-correct and order-independent.

    Args:
        artifact: The steering artifact, given as a `SteeringVector`, a per-layer directions
            mapping (`Mapping[int, Tensor]`, each `[K, H]` or `[H]` treated as `K=1`), or an
            `ArtifactSource` (unbound until `bind(ctx)`). Required.
        alpha: Ablation strength in `[0, 1]`. `1.0` = full removal (default); `< 1.0` = partial.

    Reference:

    - "Refusal in Language Models Is Mediated by a Single Direction"
      Andy Arditi, Oscar Obeso, Aaquib Syed, Daniel Paleka, Nina Panickssery, Wes Gurnee,
      Neel Nanda
      [https://arxiv.org/abs/2406.11717](https://arxiv.org/abs/2406.11717)
    """

    wire_kind: ClassVar[str | None] = "projection"

    def __init__(
        self,
        artifact: SteeringVector | Mapping[int, torch.Tensor] | ArtifactSource,
        alpha: float = 1.0,
    ):
        self.alpha = float(alpha)
        self._source: ArtifactSource | None = None
        self.directions: dict[int, torch.Tensor] | None = None
        self._basis_cache: dict[tuple, torch.Tensor] = {}  # (layer_id, device, dtype) -> [K, H] orthonormal

        self._artifact_meta: dict | None = None
        if isinstance(artifact, ArtifactSource):
            self._source = artifact
        elif isinstance(artifact, SteeringVector):
            self.directions = artifact.directions
            self._artifact_meta = dict(artifact.meta) if artifact.meta else None
        elif isinstance(artifact, Mapping):
            self.directions = dict(artifact)
        else:
            raise TypeError(
                f"ProjectionTransform artifact must be a SteeringVector, a "
                f"Mapping[int, Tensor], or an ArtifactSource; got {type(artifact).__name__} "
                f"(did you mean alpha={artifact!r}?)."
            )

    @property
    def is_bound(self) -> bool:
        return self.directions is not None

    @property
    def artifact_meta(self) -> dict | None:
        return self._artifact_meta

    def bind(self, ctx: "TransformContext") -> "ProjectionTransform":
        if self.is_bound:
            return self
        return ProjectionTransform(ctx.resolve(self._source), alpha=self.alpha)

    @property
    def covered_layer_ids(self) -> set[int] | None:
        return set(self.directions.keys()) if self.directions is not None else None

    def to_config(self) -> tuple[dict, object | None, "BaseTransform | None"]:
        """The `(params, artifact, inner)` serialized form under the `projection` kind."""
        params = {"alpha": float(self.alpha)}
        if self.directions is not None:
            artifact = SteeringVector(
                model_type="unknown", directions=dict(self.directions), meta=dict(self._artifact_meta or {}),
            )
        else:
            artifact = self._source
        return params, artifact, None

    @classmethod
    def from_config(cls, params: dict, *, artifact=None, inner=None) -> "ProjectionTransform":
        """Rebuild a `projection` transform from its serialized form."""
        return cls(artifact, alpha=params.get("alpha", 1.0))

    def wire_plan(self) -> str | None:
        """`"projection"` for single-direction full removal; None otherwise.

        The wire kind removes a single direction's component in full, so only `K == 1`
        directions at `alpha == 1.0` serialize; subspace ablation (`K > 1`) and graded
        removal (`alpha < 1.0`) are hook-only.
        """
        if self.alpha != 1.0:
            return None
        if self.directions is not None and any(
            direction.ndim == 2 and direction.size(0) > 1 for direction in self.directions.values()
        ):
            return None
        return "projection"

    def export(self, layer_id: int) -> "WireForm | None":
        """The `projection` wire form for `layer_id`, or None when the configuration is
        hook-only (`K > 1` or `alpha != 1.0`)."""
        from ..specs import WireForm

        if self.directions is None or self.alpha != 1.0:
            return None
        direction = self.directions.get(layer_id)
        if direction is None:
            return None
        if direction.ndim == 2:
            if direction.size(0) != 1:
                return None
            direction = direction.squeeze(0)
        return WireForm(kind="projection", tensors={"vector": direction})


    def _basis(self, layer_id: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
        """Return the cached orthonormal `[K, H]` basis for a layer, computing it on first use."""
        key = (layer_id, device, dtype)
        cached = self._basis_cache.get(key)
        if cached is None:
            raw = self.directions[layer_id].to(device=device, dtype=dtype)
            if raw.ndim == 1:
                raw = raw.unsqueeze(0)  # [H] -> [1, H]

            # Gram-Schmidt -> orthonormal rows (drop near-zero / dependent rows)
            basis_rows: list[torch.Tensor] = []
            for row in raw:
                v = row.clone()
                for b in basis_rows:
                    v = v - (v @ b) * b
                n = v.norm()
                if n > 1e-8:
                    basis_rows.append(v / n)

            cached = (
                torch.stack(basis_rows, dim=0)
                if basis_rows
                else raw.new_zeros((0, raw.size(-1)))
            )
            self._basis_cache[key] = cached
        return cached

    def apply(
        self,
        hidden_states: torch.Tensor,
        *,
        layer_id: int,
        token_mask: torch.BoolTensor,
        **kwargs,
    ) -> torch.Tensor:
        """Project the learned directions out of each masked hidden state.

        Args:
            hidden_states: Shape `[B, T, H]`.
            layer_id: Which layer this is being applied at.
            token_mask: Shape `[B, T]`. True at positions to ablate.

        Returns:
            Modified hidden states, same shape as input.
        """
        self._require_bound()
        if layer_id not in self.directions:
            return hidden_states

        basis = self._basis(layer_id, hidden_states.device, hidden_states.dtype)
        if basis.size(0) == 0:
            return hidden_states

        # coefficients c[..., i] = h . d_i  ->  [B, T, K]
        coeffs = torch.einsum("bth,kh->btk", hidden_states, basis)
        # component to remove: sum_i c_i d_i  ->  [B, T, H]
        component = torch.einsum("btk,kh->bth", coeffs, basis)
        ablated = hidden_states - self.alpha * component

        return torch.where(token_mask.unsqueeze(-1), ablated, hidden_states)

_artifact_meta = None instance-attribute

_basis_cache = {} instance-attribute

_source = None instance-attribute

alpha = float(alpha) instance-attribute

artifact_meta property

Provenance metadata of the transform's steering artifact, or None.

Populated when the artifact was supplied or resolved as a SteeringVector carrying meta (fit fingerprints); consumers cross-check it against a serving engine's model identity. Wrappers delegate to their inner transform. The default returns None.

covered_layer_ids property

Layers this transform can act on; None = unknown (opts out of adapter validation).

Artifact transforms report set(directions.keys()) when bound and None when unbound; wrappers delegate to their inner transform.

directions = None instance-attribute

is_bound property

True when the transform can be applied as-is (artifact concrete).

Default True: transforms that take no steering artifact are always bound.

is_modifier = False class-attribute

source property

The unresolved ArtifactSource this transform carries, or None when concrete.

wire_kind = 'projection' class-attribute

__init__(artifact, alpha=1.0)

Source code in steerability/algorithms/state_control/common/transforms/projection.py

def __init__(
    self,
    artifact: SteeringVector | Mapping[int, torch.Tensor] | ArtifactSource,
    alpha: float = 1.0,
):
    self.alpha = float(alpha)
    self._source: ArtifactSource | None = None
    self.directions: dict[int, torch.Tensor] | None = None
    self._basis_cache: dict[tuple, torch.Tensor] = {}  # (layer_id, device, dtype) -> [K, H] orthonormal

    self._artifact_meta: dict | None = None
    if isinstance(artifact, ArtifactSource):
        self._source = artifact
    elif isinstance(artifact, SteeringVector):
        self.directions = artifact.directions
        self._artifact_meta = dict(artifact.meta) if artifact.meta else None
    elif isinstance(artifact, Mapping):
        self.directions = dict(artifact)
    else:
        raise TypeError(
            f"ProjectionTransform artifact must be a SteeringVector, a "
            f"Mapping[int, Tensor], or an ArtifactSource; got {type(artifact).__name__} "
            f"(did you mean alpha={artifact!r}?)."
        )

_basis(layer_id, device, dtype)

Return the cached orthonormal [K, H] basis for a layer, computing it on first use.

Source code in steerability/algorithms/state_control/common/transforms/projection.py

def _basis(self, layer_id: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    """Return the cached orthonormal `[K, H]` basis for a layer, computing it on first use."""
    key = (layer_id, device, dtype)
    cached = self._basis_cache.get(key)
    if cached is None:
        raw = self.directions[layer_id].to(device=device, dtype=dtype)
        if raw.ndim == 1:
            raw = raw.unsqueeze(0)  # [H] -> [1, H]

        # Gram-Schmidt -> orthonormal rows (drop near-zero / dependent rows)
        basis_rows: list[torch.Tensor] = []
        for row in raw:
            v = row.clone()
            for b in basis_rows:
                v = v - (v @ b) * b
            n = v.norm()
            if n > 1e-8:
                basis_rows.append(v / n)

        cached = (
            torch.stack(basis_rows, dim=0)
            if basis_rows
            else raw.new_zeros((0, raw.size(-1)))
        )
        self._basis_cache[key] = cached
    return cached

_require_bound()

Raise a clear RuntimeError if this transform is unbound.

Called as the first line of every artifact-consuming apply.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def _require_bound(self) -> None:
    """Raise a clear `RuntimeError` if this transform is unbound.

    Called as the first line of every artifact-consuming `apply`.
    """
    if not self.is_bound:
        raise RuntimeError(
            f"{type(self).__name__} was constructed with an ArtifactSource and is unbound; pass "
            f"it through a control that resolves and binds it during steer() (ActivationAdapter, "
            f"or CAST via behavior_transform)."
        )

apply(hidden_states, *, layer_id, token_mask, **kwargs)

Project the learned directions out of each masked hidden state.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Shape [B, T, H].	required
layer_id	int	Which layer this is being applied at.	required
token_mask	BoolTensor	Shape [B, T]. True at positions to ablate.	required

Returns:

Type	Description
Tensor	Modified hidden states, same shape as input.

Source code in steerability/algorithms/state_control/common/transforms/projection.py

def apply(
    self,
    hidden_states: torch.Tensor,
    *,
    layer_id: int,
    token_mask: torch.BoolTensor,
    **kwargs,
) -> torch.Tensor:
    """Project the learned directions out of each masked hidden state.

    Args:
        hidden_states: Shape `[B, T, H]`.
        layer_id: Which layer this is being applied at.
        token_mask: Shape `[B, T]`. True at positions to ablate.

    Returns:
        Modified hidden states, same shape as input.
    """
    self._require_bound()
    if layer_id not in self.directions:
        return hidden_states

    basis = self._basis(layer_id, hidden_states.device, hidden_states.dtype)
    if basis.size(0) == 0:
        return hidden_states

    # coefficients c[..., i] = h . d_i  ->  [B, T, K]
    coeffs = torch.einsum("bth,kh->btk", hidden_states, basis)
    # component to remove: sum_i c_i d_i  ->  [B, T, H]
    component = torch.einsum("btk,kh->bth", coeffs, basis)
    ablated = hidden_states - self.alpha * component

    return torch.where(token_mask.unsqueeze(-1), ablated, hidden_states)

bind(ctx)

Return a fully-bound transform for this context.

Contract:

- MUST NOT mutate `self`; instances and sources are shared across adapters and
  `ControlSpec` grid points, whose params objects are reused per point.
- Returns `self` when already bound (idempotent).
- When source-carrying, returns a NEW instance of the same class constructed with
  `ctx.resolve(self._source)` and all hyperparameters copied, built by fresh
  construction rather than `copy.copy`.

Default: returns self.

Source code in steerability/algorithms/state_control/common/transforms/projection.py

def bind(self, ctx: "TransformContext") -> "ProjectionTransform":
    if self.is_bound:
        return self
    return ProjectionTransform(ctx.resolve(self._source), alpha=self.alpha)

export(layer_id)

The projection wire form for layer_id, or None when the configuration is hook-only (K > 1 or alpha != 1.0).

Source code in steerability/algorithms/state_control/common/transforms/projection.py

def export(self, layer_id: int) -> "WireForm | None":
    """The `projection` wire form for `layer_id`, or None when the configuration is
    hook-only (`K > 1` or `alpha != 1.0`)."""
    from ..specs import WireForm

    if self.directions is None or self.alpha != 1.0:
        return None
    direction = self.directions.get(layer_id)
    if direction is None:
        return None
    if direction.ndim == 2:
        if direction.size(0) != 1:
            return None
        direction = direction.squeeze(0)
    return WireForm(kind="projection", tensors={"vector": direction})

export_modifier(layer_id)

The wrapper's wire modifier form at layer_id, or None to contribute no modifier at that layer.

Meaningful only on wrapper transforms (is_modifier True); kind-level inexpressibility is reported by modifier_wire_kind. The default returns None.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def export_modifier(self, layer_id: int) -> "WireForm | None":
    """The wrapper's wire modifier form at `layer_id`, or None to contribute no modifier
    at that layer.

    Meaningful only on wrapper transforms (`is_modifier` True); kind-level
    inexpressibility is reported by `modifier_wire_kind`. The default returns None.
    """
    return None

from_config(params, *, artifact=None, inner=None) classmethod

Rebuild a projection transform from its serialized form.

Source code in steerability/algorithms/state_control/common/transforms/projection.py

@classmethod
def from_config(cls, params: dict, *, artifact=None, inner=None) -> "ProjectionTransform":
    """Rebuild a `projection` transform from its serialized form."""
    return cls(artifact, alpha=params.get("alpha", 1.0))

modifier_wire_kind(core_kind)

The wire modifier kind this wrapper contributes over a core transform kind, or None when the combination has no wire form.

Meaningful only on wrapper transforms (is_modifier True). The default returns the class wire_kind.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def modifier_wire_kind(self, core_kind: str) -> str | None:
    """The wire modifier kind this wrapper contributes over a core transform kind, or
    None when the combination has no wire form.

    Meaningful only on wrapper transforms (`is_modifier` True). The default returns the
    class `wire_kind`.
    """
    return type(self).wire_kind

to_config()

The (params, artifact, inner) serialized form under the projection kind.

Source code in steerability/algorithms/state_control/common/transforms/projection.py

def to_config(self) -> tuple[dict, object | None, "BaseTransform | None"]:
    """The `(params, artifact, inner)` serialized form under the `projection` kind."""
    params = {"alpha": float(self.alpha)}
    if self.directions is not None:
        artifact = SteeringVector(
            model_type="unknown", directions=dict(self.directions), meta=dict(self._artifact_meta or {}),
        )
    else:
        artifact = self._source
    return params, artifact, None

wire_plan()

"projection" for single-direction full removal; None otherwise.

The wire kind removes a single direction's component in full, so only K == 1 directions at alpha == 1.0 serialize; subspace ablation (K > 1) and graded removal (alpha < 1.0) are hook-only.

Source code in steerability/algorithms/state_control/common/transforms/projection.py

def wire_plan(self) -> str | None:
    """`"projection"` for single-direction full removal; None otherwise.

    The wire kind removes a single direction's component in full, so only `K == 1`
    directions at `alpha == 1.0` serialize; subspace ablation (`K > 1`) and graded
    removal (`alpha < 1.0`) are hook-only.
    """
    if self.alpha != 1.0:
        return None
    if self.directions is not None and any(
        direction.ndim == 2 and direction.size(0) > 1 for direction in self.directions.values()
    ):
        return None
    return "projection"

rotation

Angular (rotational) activation steering transform.

RotationMode = Literal['target', 'offset'] module-attribute

RotationTransform

Bases: BaseTransform

Rotates hidden states within a per-layer 2D steering plane.

Given an orthonormal basis pair (b1, b2) per layer, decompose each hidden state into its in-plane coordinates (c1, c2) = (h·b1, h·b2) and an out-of-plane remainder h_perp = h - c1 b1 - c2 b2, rotate the in-plane part, then reconstruct. The out-of-plane component is untouched; because a 2D rotation is orthogonal the in-plane magnitude (hence the full activation norm) is preserved by construction.

Modes:

• "target": rotate the in-plane component TO the absolute angle angle measured from b1 (paper Eq. 2). Sweeping over [0, 2*pi) traces the behavior circle.
• "offset": rotate the in-plane component BY angle (the same rotation for every token); recovers vector-addition / ablation as special cases.

The stored basis need not be orthonormal; it is re-orthonormalized via Gram-Schmidt and cached per (layer_id, device, dtype). Only positions selected by token_mask are modified; at every other position the input is returned unchanged.

Parameters:

Name	Type	Description	Default
artifact	SteeringVector | Mapping[int, Tensor] | ArtifactSource	The steering artifact whose per-layer directions are [2, H] (row 0 = feature axis, row 1 = companion axis), given as a SteeringVector, a per-layer directions mapping, or an ArtifactSource (unbound until bind(ctx)). Required.	required
angle	float	Rotation angle in radians. Interpreted as an absolute target angle in "target" mode and as a signed offset in "offset" mode.	0.0
mode	RotationMode	Either "target" or "offset".	'target'

Raises:

Type	Description
ValueError	If any layer's directions are not shaped [2, H], or if mode is invalid.

Reference:

• "Angular Steering: Behavior Control via Rotation in Activation Space" Hieu M. Vu, Tan M. Nguyen https://arxiv.org/abs/2510.26243

Source code in steerability/algorithms/state_control/common/transforms/rotation.py

class RotationTransform(BaseTransform):
    """Rotates hidden states within a per-layer 2D steering plane.

    Given an orthonormal basis pair `(b1, b2)` per layer, decompose each hidden state into its
    in-plane coordinates `(c1, c2) = (h·b1, h·b2)` and an out-of-plane remainder
    `h_perp = h - c1 b1 - c2 b2`, rotate the in-plane part, then reconstruct. The out-of-plane
    component is untouched; because a 2D rotation is orthogonal the in-plane magnitude (hence the
    full activation norm) is preserved by construction.

    Modes:

    - `"target"`: rotate the in-plane component TO the absolute angle `angle` measured from `b1`
        (paper Eq. 2). Sweeping over `[0, 2*pi)` traces the behavior circle.
    - `"offset"`: rotate the in-plane component BY `angle` (the same rotation for every token);
        recovers vector-addition / ablation as special cases.

    The stored basis need not be orthonormal; it is re-orthonormalized via Gram-Schmidt and cached
    per `(layer_id, device, dtype)`. Only positions selected by `token_mask` are modified; at every
    other position the input is returned unchanged.

    Args:
        artifact: The steering artifact whose per-layer directions are `[2, H]` (row 0 = feature
            axis, row 1 = companion axis), given as a `SteeringVector`, a per-layer directions
            mapping, or an `ArtifactSource` (unbound until `bind(ctx)`). Required.
        angle: Rotation angle in radians. Interpreted as an absolute target angle in `"target"`
            mode and as a signed offset in `"offset"` mode.
        mode: Either `"target"` or `"offset"`.

    Raises:
        ValueError: If any layer's directions are not shaped `[2, H]`, or if `mode` is invalid.

    Reference:

    - "Angular Steering: Behavior Control via Rotation in Activation Space"
      Hieu M. Vu, Tan M. Nguyen
      [https://arxiv.org/abs/2510.26243](https://arxiv.org/abs/2510.26243)
    """

    wire_kind: ClassVar[str | None] = "rotation"

    def __init__(
        self,
        artifact: SteeringVector | Mapping[int, torch.Tensor] | ArtifactSource,
        angle: float = 0.0,
        mode: RotationMode = "target",
    ):
        if mode not in ("target", "offset"):
            raise ValueError(f"mode must be 'target' or 'offset'; got {mode!r}.")

        self.angle = float(angle)
        self.mode = mode
        self._source: ArtifactSource | None = None
        self.steering_vector: SteeringVector | None = None
        self._basis_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {}

        if isinstance(artifact, ArtifactSource):
            self._source = artifact
        elif isinstance(artifact, SteeringVector):
            self.steering_vector = artifact
            self._validate_artifact()
        elif isinstance(artifact, Mapping):
            self.steering_vector = SteeringVector(model_type="unknown", directions=dict(artifact))
            self._validate_artifact()
        else:
            raise TypeError(
                f"RotationTransform artifact must be a SteeringVector, a Mapping[int, Tensor], or an "
                f"ArtifactSource; got {type(artifact).__name__} (did you mean angle=?)."
            )

    def _validate_artifact(self) -> None:
        """Validate the per-layer directions are `[2, H]` basis pairs."""
        for layer_id, dirs in self.steering_vector.directions.items():
            if dirs.ndim != 2 or dirs.size(0) != 2:
                raise ValueError(
                    f"RotationTransform expects [2, H] directions (basis pair); "
                    f"got shape {tuple(dirs.shape)} at layer {layer_id}."
                )

    @property
    def is_bound(self) -> bool:
        return self.steering_vector is not None

    @property
    def artifact_meta(self) -> dict | None:
        if self.steering_vector is not None and self.steering_vector.meta:
            return dict(self.steering_vector.meta)
        return None

    def bind(self, ctx: "TransformContext") -> "RotationTransform":
        if self.is_bound:
            return self
        return RotationTransform(ctx.resolve(self._source), angle=self.angle, mode=self.mode)

    @property
    def covered_layer_ids(self) -> set[int] | None:
        return set(self.steering_vector.directions.keys()) if self.steering_vector is not None else None

    def to_config(self) -> tuple[dict, object | None, "BaseTransform | None"]:
        """The `(params, artifact, inner)` serialized form under the `rotation` kind."""
        params = {"angle": float(self.angle), "mode": self.mode}
        artifact = self.steering_vector if self.steering_vector is not None else self._source
        return params, artifact, None

    @classmethod
    def from_config(cls, params: dict, *, artifact=None, inner=None) -> "RotationTransform":
        """Rebuild a `rotation` transform from its serialized form."""
        return cls(artifact, angle=params.get("angle", 0.0), mode=params.get("mode", "target"))

    def export(self, layer_id: int) -> "WireForm | None":
        """The `rotation` wire form for `layer_id` (angle, mode, and the `[2, H]` basis)."""
        from ..specs import WireForm

        if self.steering_vector is None:
            return None
        basis = self.steering_vector.directions.get(layer_id)
        if basis is None:
            return None
        return WireForm(
            kind="rotation",
            params={"angle": float(self.angle), "mode": self.mode},
            tensors={"basis": basis},
        )


    def _basis(self, layer_id: int, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
        """Return the cached orthonormal `(b1, b2)` for a layer, computing it on first use."""
        key = (layer_id, device, dtype)
        cached = self._basis_cache.get(key)
        if cached is None:
            raw = self.steering_vector.directions[layer_id].to(device=device, dtype=dtype)
            b1 = raw[0] / (raw[0].norm() + 1e-8)
            b2 = raw[1] - (raw[1] @ b1) * b1  # Gram-Schmidt
            b2 = b2 / (b2.norm() + 1e-8)
            cached = (b1, b2)
            self._basis_cache[key] = cached
        return cached

    def apply(
        self,
        hidden_states: torch.Tensor,
        *,
        layer_id: int,
        token_mask: torch.BoolTensor,
        **kwargs,
    ) -> torch.Tensor:
        """Rotate the in-plane component of each masked hidden state.

        Args:
            hidden_states: Shape `[B, T, H]`.
            layer_id: Which layer this is being applied at.
            token_mask: Shape `[B, T]`. True at positions to rotate.

        Returns:
            Modified hidden states, same shape as input.
        """
        self._require_bound()
        if layer_id not in self.steering_vector.directions:
            return hidden_states

        b1, b2 = self._basis(layer_id, hidden_states.device, hidden_states.dtype)

        c1 = hidden_states @ b1  # [B, T]
        c2 = hidden_states @ b2  # [B, T]

        if self.mode == "target":
            delta = self.angle - torch.atan2(c2, c1)  # per-token, to hit the absolute angle
            cos_d, sin_d = torch.cos(delta), torch.sin(delta)
        else:
            cos_d = hidden_states.new_tensor(math.cos(self.angle))
            sin_d = hidden_states.new_tensor(math.sin(self.angle))

        c1_new = cos_d * c1 - sin_d * c2
        c2_new = sin_d * c1 + cos_d * c2

        # h' = h + (c1_new - c1) b1 + (c2_new - c2) b2
        delta_c1 = (c1_new - c1).unsqueeze(-1)
        delta_c2 = (c2_new - c2).unsqueeze(-1)
        rotated = hidden_states + delta_c1 * b1 + delta_c2 * b2

        return torch.where(token_mask.unsqueeze(-1), rotated, hidden_states)

_basis_cache = {} instance-attribute

_source = None instance-attribute

angle = float(angle) instance-attribute

artifact_meta property

Provenance metadata of the transform's steering artifact, or None.

Populated when the artifact was supplied or resolved as a SteeringVector carrying meta (fit fingerprints); consumers cross-check it against a serving engine's model identity. Wrappers delegate to their inner transform. The default returns None.

covered_layer_ids property

Layers this transform can act on; None = unknown (opts out of adapter validation).

Artifact transforms report set(directions.keys()) when bound and None when unbound; wrappers delegate to their inner transform.

is_bound property

True when the transform can be applied as-is (artifact concrete).

Default True: transforms that take no steering artifact are always bound.

is_modifier = False class-attribute

mode = mode instance-attribute

source property

The unresolved ArtifactSource this transform carries, or None when concrete.

steering_vector = None instance-attribute

wire_kind = 'rotation' class-attribute

__init__(artifact, angle=0.0, mode='target')

Source code in steerability/algorithms/state_control/common/transforms/rotation.py

def __init__(
    self,
    artifact: SteeringVector | Mapping[int, torch.Tensor] | ArtifactSource,
    angle: float = 0.0,
    mode: RotationMode = "target",
):
    if mode not in ("target", "offset"):
        raise ValueError(f"mode must be 'target' or 'offset'; got {mode!r}.")

    self.angle = float(angle)
    self.mode = mode
    self._source: ArtifactSource | None = None
    self.steering_vector: SteeringVector | None = None
    self._basis_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {}

    if isinstance(artifact, ArtifactSource):
        self._source = artifact
    elif isinstance(artifact, SteeringVector):
        self.steering_vector = artifact
        self._validate_artifact()
    elif isinstance(artifact, Mapping):
        self.steering_vector = SteeringVector(model_type="unknown", directions=dict(artifact))
        self._validate_artifact()
    else:
        raise TypeError(
            f"RotationTransform artifact must be a SteeringVector, a Mapping[int, Tensor], or an "
            f"ArtifactSource; got {type(artifact).__name__} (did you mean angle=?)."
        )

_basis(layer_id, device, dtype)

Return the cached orthonormal (b1, b2) for a layer, computing it on first use.

Source code in steerability/algorithms/state_control/common/transforms/rotation.py

def _basis(self, layer_id: int, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
    """Return the cached orthonormal `(b1, b2)` for a layer, computing it on first use."""
    key = (layer_id, device, dtype)
    cached = self._basis_cache.get(key)
    if cached is None:
        raw = self.steering_vector.directions[layer_id].to(device=device, dtype=dtype)
        b1 = raw[0] / (raw[0].norm() + 1e-8)
        b2 = raw[1] - (raw[1] @ b1) * b1  # Gram-Schmidt
        b2 = b2 / (b2.norm() + 1e-8)
        cached = (b1, b2)
        self._basis_cache[key] = cached
    return cached

_require_bound()

Raise a clear RuntimeError if this transform is unbound.

Called as the first line of every artifact-consuming apply.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def _require_bound(self) -> None:
    """Raise a clear `RuntimeError` if this transform is unbound.

    Called as the first line of every artifact-consuming `apply`.
    """
    if not self.is_bound:
        raise RuntimeError(
            f"{type(self).__name__} was constructed with an ArtifactSource and is unbound; pass "
            f"it through a control that resolves and binds it during steer() (ActivationAdapter, "
            f"or CAST via behavior_transform)."
        )

_validate_artifact()

Validate the per-layer directions are [2, H] basis pairs.

Source code in steerability/algorithms/state_control/common/transforms/rotation.py

def _validate_artifact(self) -> None:
    """Validate the per-layer directions are `[2, H]` basis pairs."""
    for layer_id, dirs in self.steering_vector.directions.items():
        if dirs.ndim != 2 or dirs.size(0) != 2:
            raise ValueError(
                f"RotationTransform expects [2, H] directions (basis pair); "
                f"got shape {tuple(dirs.shape)} at layer {layer_id}."
            )

apply(hidden_states, *, layer_id, token_mask, **kwargs)

Rotate the in-plane component of each masked hidden state.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Shape [B, T, H].	required
layer_id	int	Which layer this is being applied at.	required
token_mask	BoolTensor	Shape [B, T]. True at positions to rotate.	required

Returns:

Type	Description
Tensor	Modified hidden states, same shape as input.

Source code in steerability/algorithms/state_control/common/transforms/rotation.py

def apply(
    self,
    hidden_states: torch.Tensor,
    *,
    layer_id: int,
    token_mask: torch.BoolTensor,
    **kwargs,
) -> torch.Tensor:
    """Rotate the in-plane component of each masked hidden state.

    Args:
        hidden_states: Shape `[B, T, H]`.
        layer_id: Which layer this is being applied at.
        token_mask: Shape `[B, T]`. True at positions to rotate.

    Returns:
        Modified hidden states, same shape as input.
    """
    self._require_bound()
    if layer_id not in self.steering_vector.directions:
        return hidden_states

    b1, b2 = self._basis(layer_id, hidden_states.device, hidden_states.dtype)

    c1 = hidden_states @ b1  # [B, T]
    c2 = hidden_states @ b2  # [B, T]

    if self.mode == "target":
        delta = self.angle - torch.atan2(c2, c1)  # per-token, to hit the absolute angle
        cos_d, sin_d = torch.cos(delta), torch.sin(delta)
    else:
        cos_d = hidden_states.new_tensor(math.cos(self.angle))
        sin_d = hidden_states.new_tensor(math.sin(self.angle))

    c1_new = cos_d * c1 - sin_d * c2
    c2_new = sin_d * c1 + cos_d * c2

    # h' = h + (c1_new - c1) b1 + (c2_new - c2) b2
    delta_c1 = (c1_new - c1).unsqueeze(-1)
    delta_c2 = (c2_new - c2).unsqueeze(-1)
    rotated = hidden_states + delta_c1 * b1 + delta_c2 * b2

    return torch.where(token_mask.unsqueeze(-1), rotated, hidden_states)

bind(ctx)

Return a fully-bound transform for this context.

Contract:

- MUST NOT mutate `self`; instances and sources are shared across adapters and
  `ControlSpec` grid points, whose params objects are reused per point.
- Returns `self` when already bound (idempotent).
- When source-carrying, returns a NEW instance of the same class constructed with
  `ctx.resolve(self._source)` and all hyperparameters copied, built by fresh
  construction rather than `copy.copy`.

Default: returns self.

Source code in steerability/algorithms/state_control/common/transforms/rotation.py

def bind(self, ctx: "TransformContext") -> "RotationTransform":
    if self.is_bound:
        return self
    return RotationTransform(ctx.resolve(self._source), angle=self.angle, mode=self.mode)

export(layer_id)

The rotation wire form for layer_id (angle, mode, and the [2, H] basis).

Source code in steerability/algorithms/state_control/common/transforms/rotation.py

def export(self, layer_id: int) -> "WireForm | None":
    """The `rotation` wire form for `layer_id` (angle, mode, and the `[2, H]` basis)."""
    from ..specs import WireForm

    if self.steering_vector is None:
        return None
    basis = self.steering_vector.directions.get(layer_id)
    if basis is None:
        return None
    return WireForm(
        kind="rotation",
        params={"angle": float(self.angle), "mode": self.mode},
        tensors={"basis": basis},
    )

export_modifier(layer_id)

The wrapper's wire modifier form at layer_id, or None to contribute no modifier at that layer.

Meaningful only on wrapper transforms (is_modifier True); kind-level inexpressibility is reported by modifier_wire_kind. The default returns None.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def export_modifier(self, layer_id: int) -> "WireForm | None":
    """The wrapper's wire modifier form at `layer_id`, or None to contribute no modifier
    at that layer.

    Meaningful only on wrapper transforms (`is_modifier` True); kind-level
    inexpressibility is reported by `modifier_wire_kind`. The default returns None.
    """
    return None

from_config(params, *, artifact=None, inner=None) classmethod

Rebuild a rotation transform from its serialized form.

Source code in steerability/algorithms/state_control/common/transforms/rotation.py

@classmethod
def from_config(cls, params: dict, *, artifact=None, inner=None) -> "RotationTransform":
    """Rebuild a `rotation` transform from its serialized form."""
    return cls(artifact, angle=params.get("angle", 0.0), mode=params.get("mode", "target"))

modifier_wire_kind(core_kind)

The wire modifier kind this wrapper contributes over a core transform kind, or None when the combination has no wire form.

Meaningful only on wrapper transforms (is_modifier True). The default returns the class wire_kind.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def modifier_wire_kind(self, core_kind: str) -> str | None:
    """The wire modifier kind this wrapper contributes over a core transform kind, or
    None when the combination has no wire form.

    Meaningful only on wrapper transforms (`is_modifier` True). The default returns the
    class `wire_kind`.
    """
    return type(self).wire_kind

to_config()

The (params, artifact, inner) serialized form under the rotation kind.

Source code in steerability/algorithms/state_control/common/transforms/rotation.py

def to_config(self) -> tuple[dict, object | None, "BaseTransform | None"]:
    """The `(params, artifact, inner)` serialized form under the `rotation` kind."""
    params = {"angle": float(self.angle), "mode": self.mode}
    artifact = self.steering_vector if self.steering_vector is not None else self._source
    return params, artifact, None

wire_plan()

The wire transform kind this configuration serializes to, or None when hook-only.

Readable on the unbound form: a source-carrying transform consults its source's declared shape rather than resolving it. The default returns the class wire_kind.

Source code in steerability/algorithms/state_control/common/transforms/base.py

def wire_plan(self) -> str | None:
    """The wire transform kind this configuration serializes to, or None when hook-only.

    Readable on the unbound form: a source-carrying transform consults its source's
    declared shape rather than resolving it. The default returns the class `wire_kind`.
    """
    return type(self).wire_kind