interpretune.analysis.backends#

Model backends for analysis operations.

Provides the ModelBackend protocol, shared intervention helpers, and backend implementations for different model execution frameworks (TransformerLens, nnsight, etc.).

Functions

`apply_feature_score_sign_filter`(scores[, ...])	Return a boolean mask selecting feature scores with the requested sign.
`apply_feature_selection_filter`(...)	Return a boolean mask (length N) selecting rows of active_features that match spec.
`apply_intervention_to_last_token`(value, ...)	Apply one intervention spec to the last-token slice of an activation tensor.
`build_intervention_dict`(interventions, ...)	Canonicalize raw intervention inputs into a resolved `InterventionDict`.
`expand_intervention_patterns`(patterns, ...)	Expand raw hook-name patterns to ordered lists of concrete hook names.
`get_analysis_backend`(module)
`get_intervention_target_shape`(activation)	Return the per-example shape targeted by last-token interventions.
`get_model_backend`(module)	Return the module's model backend while avoiding mock-created private attrs.
`get_module_capabilities`(module)	Aggregate execution and analysis capabilities exposed by a module.
`normalize_backend_capability`(capability)	Normalize capability-like values to the local execution or analysis capability enums.
`require_analysis_backend`(module)	Return the module's analysis backend or raise if it is unavailable.
`resolve_interventions`(*, analysis_batch, ...)	Resolve explicit or shorthand intervention inputs into a standardized payload mapping.

Classes

`AnalysisBackend`(args, *kwargs)	Protocol defining analysis-adapter functionality layered above model execution backends.
`AnalysisBackendCapability`(*values)	Capabilities exposed by analysis adapters/backends rather than model execution backends.
`BackendCapability`(*values)	Capabilities that a model backend may support.
`FeatureSelectionSpec`([layers, positions, ...])	Pre-filter specification for `extract_top_features_impl()`.
`InterventionDict`(hook_map)	Canonical mapping from resolved hook names to intervention specs.
`InterventionSpec`(intervention_tensor[, ...])	Specification for a single hook-point intervention.
`ModelBackend`(args, *kwargs)	Protocol defining the interface for model execution backends.
`ModuleCapabilities`(model, analysis)	Execution and analysis capabilities exposed by a module.

class interpretune.analysis.backends.AnalysisBackend(*args, **kwargs)[source]#

Protocol defining analysis-adapter functionality layered above model execution backends.

supports(capability)[source]#

Check whether this backend supports a given analysis capability.

Return type:: bool
Parameters:: capability (AnalysisBackendCapability)

property capabilities: frozenset[AnalysisBackendCapability]#: Return the set of analysis capabilities this backend supports.

class interpretune.analysis.backends.AnalysisBackendCapability(*values)[source]#

Capabilities exposed by analysis adapters/backends rather than model execution backends.

ATTRIBUTION_GRAPH = 'attribution_graph'#: Module exposes attribution graph analysis support via an attached analysis backend.

FEATURE_INTERVENTION = 'feature_intervention'#: Module exposes feature intervention support via an attached analysis backend.

class interpretune.analysis.backends.BackendCapability(*values)[source]#

Capabilities that a model backend may support.

Ops and the dispatcher can query backend.capabilities to check support before calling optional methods. Backends that do not support a capability should fall back to a simpler code path (e.g., looping instead of batching).

BATCHED_HOOKS = 'batched_hooks'#: Backend can run multiple forward passes with different hook configs in a single batched execution (e.g., NNsight multi-invoke within one trace).

GRADIENTS = 'gradients'#: Backend supports forward + backward with gradient caching.

class interpretune.analysis.backends.FeatureSelectionSpec(layers=<factory>, positions=<factory>, feature_ids=<factory>, layer_slice=None, position_slice=None, triples=<factory>, layer_feature_pairs=<factory>, activation_overrides=<factory>, score_source=None, score_sign='any', rank_by_abs=False)[source]#

Pre-filter specification for extract_top_features_impl().

All criteria use OR semantics: a feature row (layer, position, feature_id) passes the filter if it matches any of the non-empty criteria.

Numeric slice notation is supported for layers and positions — pass a Python slice object alongside (or instead of) explicit int lists. The slice is applied as a numeric range over the observed values in active_features, so slice(10, None) means “layer >= 10” and slice(0, 10) means “position >= 0 and < 10”.

Parameters:

layers (list[int])
positions (list[int])
feature_ids (list[int])
layer_slice (slice | None)
position_slice (slice | None)
triples (list[tuple[int, int, int]])
layer_feature_pairs (list[tuple[int, int]])
activation_overrides (dict[tuple[int, int], float])
score_source (str | None)
score_sign (str)
rank_by_abs (bool)

layers#: Explicit layer indices to include.

positions#: Explicit token-position indices to include.

feature_ids#: Explicit feature-ID values to include.

layer_slice#: A slice expanded over observed layer values.

position_slice#: A slice expanded over observed position values.

triples#: Exact (layer, position, feature_id) tuples to include.

layer_feature_pairs#: Exact (layer, feature_id) pairs to include across any position.

activation_overrides#: Optional override activation values keyed by (layer, feature_id).

score_source#: Optional analysis-batch field name or alias to use for feature ranking. Supported aliases include "influence", "signed_influence", and the planned backward-pass "gradient" / "logit_diff_gradient" channel for gradients of selected feature activations with respect to a target logit difference.

score_sign#: Optional sign filter for score values: "any", "positive", or "negative".

rank_by_abs#: If true, rank by absolute score magnitude while preserving the original signed score values.

class interpretune.analysis.backends.InterventionDict(hook_map)[source]#

Canonical mapping from resolved hook names to intervention specs.

Keys are concrete hook-point names with wildcards already expanded. Values are ordered tuples of intervention specs to apply sequentially at that hook.

Parameters:: hook_map (dict[str, tuple[InterventionSpec, ...]])

items() → a set-like object providing a view on D's items[source]#

keys() → a set-like object providing a view on D's keys[source]#

values() → an object providing a view on D's values[source]#

class interpretune.analysis.backends.InterventionSpec(intervention_tensor, mode='replace', scale_factor=1.0, use_intervention_tensor_as_basis=True)[source]#

Specification for a single hook-point intervention.

Parameters:

intervention_tensor (Tensor)
mode (str)
scale_factor (float)
use_intervention_tensor_as_basis (bool)

intervention_tensor#: Intervention tensor with any shape broadcast-compatible with the targeted hook-point slice at the intervention position. This is not restricted to (d_model,) and may instead match higher-rank activations such as (n_heads, d_head) or latent feature activations.

mode#: "replace" overwrites the activation at the target position with the tensor; "add" adds intervention_tensor * scale_factor to the activation; "project" replaces the activation with a projection result. By default, the current hook input is projected onto the span of intervention_tensor, so the intervention tensor acts as the projection basis. When use_intervention_tensor_as_basis is False, the direction is reversed and intervention_tensor is projected onto the span of the current hook input.

scale_factor#: Scalar multiplier applied to intervention_tensor before the intervention (not used in "replace" mode and applied to the projected activation in "project" mode).

use_intervention_tensor_as_basis#: Controls which vector defines the projection basis in "project" mode. True means project the current hook input onto the span of intervention_tensor. False means project intervention_tensor onto the span of the current hook input instead.

Create new instance of InterventionSpec(intervention_tensor, mode, scale_factor, use_intervention_tensor_as_basis)

intervention_tensor: Tensor#: Alias for field number 0

mode: str#: Alias for field number 1

scale_factor: float#: Alias for field number 2

use_intervention_tensor_as_basis: bool#: Alias for field number 3

class interpretune.analysis.backends.ModelBackend(*args, **kwargs)[source]#

Protocol defining the interface for model execution backends.

Each backend wraps a specific framework’s model execution API (e.g., TransformerLens hook-based execution, nnsight trace-based execution) behind a uniform interface used by analysis op implementations.

Note

hook=True evaluation

NNsight’s hook=True parameter (on tracer.invoke()) enables .output / .input access on auxiliary modules like SAEs. Our current architecture calls sae.encode() / sae.decode() explicitly within the trace, giving direct proxy access to feature activations. If SAEs were registered as model sub-modules, hook=True could replace explicit encode/decode calls, but the current external latent_model_handles design makes hook=True unnecessary. Adding hook=True would require architectural changes to how SAEs are attached and is best evaluated in a future session.

fwd(model, batch)[source]#

Run a minimal forward pass and return logits.

Each backend handles any necessary batch-key mapping (e.g., the NNsight backend wraps the call in a trace context so that LanguageModel._prepare_input correctly routes input → input_ids for HuggingFace models).

Parameters:

model (Any) – The model to run.
batch (dict[str, Any]) – Input batch dict.

Return type:

Tensor

Returns:

Model output logits tensor.

fwd_w_cache(model, batch, names_filter)[source]#

Run a forward pass with activation caching but without latent model hooks.

Parameters:

model (Any) – The model to run.
batch (dict[str, Any]) – Input batch dict.
names_filter (Union[Callable[[str], bool], Sequence[str], str, None]) – Filter specifying which hook activations to cache.

Return type:

tuple[Tensor, Any]

Returns:

Tuple of (logits, activation_cache).

fwd_w_cache_and_latent_models(model, batch, latent_model_handles, names_filter)[source]#

Run a forward pass with activation caching and latent model hooks.

Parameters:

model (Any) – The model to run (e.g., HookedSAETransformer, SAETransformerBridge).
batch (dict[str, Any]) – Input batch dict (unpacked as **batch for the model call).
latent_model_handles (list[Any]) – Latent model handles (e.g., SAE objects) to attach.
names_filter (Union[Callable[[str], bool], Sequence[str], str, None]) – Filter specifying which hook activations to cache.

Return type:

tuple[Tensor, Any]

Returns:

Tuple of (logits, activation_cache).

fwd_w_grads_and_latent_models(model, batch, latent_model_handles, fwd_hooks, bwd_hooks, backward_fn)[source]#

Run forward + backward with latent model hooks and gradient caching.

The backend owns the entire forward+backward execution flow. This enables both eager execution (TransformerLens) and deferred/traced execution (NNsight) to use the same op-level code.

The backward_fn closure is provided by the analysis op and computes a scalar metric from raw model logits. The backend calls backward_fn(logits) to obtain the scalar, then runs .backward() on it (eager for TL, deferred via with scalar.backward(): for NNsight).

Forward and backward cache hooks (fwd_hooks, bwd_hooks) are structured as [(names_filter, cache_fn), ...] and are invoked by the backend to populate analysis_cfg.cache_dict. For TL, hooks fire during execution. For NNsight, the backend calls them after the trace completes with materialized tensors.

Parameters:

model (Any) – The model to run.
batch (dict[str, Any]) – Input batch dict (unpacked as **batch for the model call).
latent_model_handles (list[Any]) – Latent model handles (e.g., SAE objects) to attach.
fwd_hooks (list[tuple[Any, Any]]) – Forward cache hooks [(names_filter, cache_fn), ...].
bwd_hooks (list[tuple[Any, Any]]) – Backward cache hooks [(names_filter, cache_fn), ...].
backward_fn (Callable[[Tensor], Tensor]) – raw_logits -> scalar. Takes the full model output logits and returns a scalar tensor to call .backward() on. Must be compatible with both real tensors (TL) and NNsight proxy objects.

Return type:

Tensor

Returns:

Raw model output logits (always a real tensor, even for NNsight).

fwd_w_hooks_and_latent_models(model, batch, latent_model_handles, fwd_hooks, clear_contexts=True)[source]#

Run a forward pass with custom forward hooks and latent model hooks.

Parameters:

model (Any) – The model to run.
batch (dict[str, Any]) – Input batch dict (unpacked as **batch for the model call).
latent_model_handles (list[Any]) – Latent model handles to attach.
fwd_hooks (list[tuple[str, Any]]) – List of (hook_name, hook_fn) tuples for forward hooks.
clear_contexts (bool) – Whether to clear hook contexts after the forward pass.

Return type:

Tensor

Returns:

Model output logits.

fwd_w_hooks_batched(model, batch, latent_model_handles, hook_configs, clear_contexts=True, configs_per_pass=None)[source]#

Run multiple forward passes with different hook configurations, batched when possible.

Each element of hook_configs is a fwd_hooks list (as passed to fwd_w_hooks_and_latent_models). Backends that support BackendCapability.BATCHED_HOOKS may batch all configs into a single execution context (e.g., NNsight multi-invoke within one trace) for efficiency. Other backends loop over configs sequentially.

configs_per_pass limits how many configs are batched per execution context. When None (default), the entire hook_configs list is batched in one context. Setting a value (e.g., 64) chunks the work to avoid OOM with very large alive-latent counts.

Note

TODO: evaluate the possibility of releasing memory after each invoke within a trace if OOMs become a problem (would require nnsight-level support).

Parameters:

model (Any) – The model to run.
batch (dict[str, Any]) – Input batch dict (unpacked as **batch for the model call).
latent_model_handles (list[Any]) – Latent model handles to attach.
hook_configs (Sequence[list[tuple[str, Any]]]) – Sequence of fwd_hooks lists, one per forward pass.
clear_contexts (bool) – Whether to clear hook contexts (for TL backend compatibility).
configs_per_pass (int | None) – Maximum number of configs to batch per execution context. None means unbounded (all configs in one trace).

Return type:

list[Tensor]

Returns:

List of logits tensors, one per element in hook_configs.

fwd_w_intervention(model, batch, interventions, latent_model_handles=None)[source]#

Run baseline + intervention forward passes using the given hook specs.

Performs two forward passes:

Baseline: captures pre-intervention logits.
Intervention: for each key in interventions, matches the key (which may contain * wildcards) against available hook names, then applies each InterventionSpec at the last sequence position according to its mode ("replace", "add", or "project").

Parameters:

model (Any) – The model to run.
batch (dict[str, Any]) – Input batch dict.
interventions (InterventionDict | Mapping[str, Any]) – Either a canonical InterventionDict keyed by concrete hook names or a raw mapping from hook-name patterns to intervention payloads. Raw payloads may be tensors, InterventionSpec instances, mapping-style specs, or sequences of those values. Patterns may use * as a glob-style wildcard.
latent_model_handles (list[Any] | None) – Optional latent model handles to enable latent-hook-aware resolution and execution.

Return type:

tuple[Any, Any]

Returns:

(pre_intervention_logits, post_intervention_logits) — both real tensors.

supports(capability)[source]#

Check whether this backend supports a given capability.

Default implementation checks capability in self.capabilities.

Return type:: bool
Parameters:: capability (BackendCapability)

wrap_activation_cache(cache_dict, model)[source]#

Wrap a raw activation dict into a backend-specific activation cache object.

For TransformerLens, wraps in ActivationCache. Other backends may return the dict as-is or wrap in their own cache type.

Parameters:

cache_dict (dict[str, Any]) – Raw dict mapping hook names to activation tensors.
model (Any) – The model instance (may be needed for cache construction).

Return type:

Any

Returns:

A cache object suitable for indexed access by hook name.

property capabilities: frozenset[BackendCapability]#

Return the set of capabilities this backend supports.

Backends must override this property to declare their capabilities. Analysis ops can check capabilities before calling optional methods.

class interpretune.analysis.backends.ModuleCapabilities(model, analysis)[source]#

Execution and analysis capabilities exposed by a module.

Parameters:

model (frozenset[BackendCapability])
analysis (frozenset[AnalysisBackendCapability])

interpretune.analysis.backends.apply_feature_score_sign_filter(scores, score_sign='any')[source]#

Return a boolean mask selecting feature scores with the requested sign.

Return type:

Tensor

Parameters:

scores (Tensor)
score_sign (str)

interpretune.analysis.backends.apply_feature_selection_filter(active_features, spec)[source]#

Return a boolean mask (length N) selecting rows of active_features that match spec.

active_features has shape (N, 3) with columns [layer, position, feature_id].

Return type:

Tensor

Parameters:

active_features (Tensor)
spec (FeatureSelectionSpec)

interpretune.analysis.backends.apply_intervention_to_last_token(value, spec, *, last_pos)[source]#

Apply one intervention spec to the last-token slice of an activation tensor.

The existing hook value is treated as the projection input and spec.intervention_tensor is treated as the projection target. In "project" mode, the target defines the default projection basis: the input is projected onto the span of the intervention tensor. When spec.use_intervention_tensor_as_basis is False, the direction is reversed and the intervention tensor is projected onto the span of the input.

Return type:

Tensor

Parameters:

value (Tensor)
spec (InterventionSpec)
last_pos (int)

interpretune.analysis.backends.build_intervention_dict(interventions, expanded_matches, hook_shapes, *, default_mode='replace', default_scale_factor=1.0)[source]#

Canonicalize raw intervention inputs into a resolved InterventionDict.

Return type:

InterventionDict

Parameters:

interventions (InterventionDict | Mapping[str, Any])
expanded_matches (Mapping[str, Sequence[str]])
hook_shapes (Mapping[str, tuple[int, ...]])
default_mode (str)
default_scale_factor (float)

interpretune.analysis.backends.expand_intervention_patterns(patterns, available_hook_map)[source]#

Expand raw hook-name patterns to ordered lists of concrete hook names.

Return type:

dict[str, list[str]]

Parameters:

patterns (Sequence[str])
available_hook_map (Mapping[str, str])

interpretune.analysis.backends.get_intervention_target_shape(activation)[source]#

Return the per-example shape targeted by last-token interventions.

Return type:: tuple[int, ...]
Parameters:: activation (Tensor)

interpretune.analysis.backends.get_module_capabilities(module)[source]#

Aggregate execution and analysis capabilities exposed by a module.

Return type:: ModuleCapabilities
Parameters:: module (Any)

interpretune.analysis.backends.normalize_backend_capability(capability)[source]#

Normalize capability-like values to the local execution or analysis capability enums.

Return type:: BackendCapability | AnalysisBackendCapability
Parameters:: capability (Any)

interpretune.analysis.backends.resolve_interventions(*, analysis_batch, resolve_field, load_json_field, kwargs=None, default_hook_qualifier='unembed.hook_in')[source]#

Resolve explicit or shorthand intervention inputs into a standardized payload mapping.

Explicit interventions or interventions_json mappings take precedence. Otherwise, shorthand op inputs are assembled into a raw intervention payload keyed by the resolved hook qualifier. Shape canonicalization into InterventionDict still happens in the backend after concrete hook shapes are known.

Return type:

InterventionDict | dict[str, Any]

Parameters:

analysis_batch (Any)
resolve_field (Callable[[str], Any])
load_json_field (Callable[[str], Any])
kwargs (Mapping[str, Any] | None)
default_hook_qualifier (str)