interpretune.analysis.ops.definitions#

Definitions of specific analysis operations.

Functions

ablate_sae_latent(sae_acts, hook[, ...])

Ablate a particular latent at a particular sequence position.

ablation_attribution_impl(module, ...[, ...])

Implementation for computing attribution values using latent ablation.

boolean_logits_to_avg_logit_diff(logits, ...)

Returns the avg logit diff on a set of prompts, with fixed s2 pos and stuff.

compute_attribution_graph_impl(module, ...)

Generate and decompose a circuit-tracer attribution graph.

concept_direction_impl(module, ...)

Compute a concept direction from latent-example rows, or fall back to token-group embeddings.

extract_concept_latent_examples_impl(module, ...)

Filter and annotate concept latent rows for downstream concept-direction aggregation.

extract_concept_latent_state_impl(module, ...)

Extract per-example latent rows from the configured cache key for downstream concept-direction ops.

extract_top_features_impl(module, ...[, top_n])

Extract the top scoring features from analysis-batch feature rows.

feature_intervention_forward_impl(module, ...)

Run circuit-tracer feature interventions against the module replacement model.

get_alive_latents_impl(module, ...)

Implementation for extracting alive latents from cache.

get_answer_indices_impl(module, ...)

Implementation for extracting answer indices from batch.

get_loss_preds_diffs(module, analysis_batch, ...)

Implementation for computing loss, predictions, and logit differences.

gradient_attribution_impl(module, ...)

Implementation for computing attribution values from gradients.

graph_node_influence_impl(module, ...)

Compute feature-node influence scores from a structured graph.

graph_prune_impl(module, analysis_batch, ...)

Prune a structured circuit-tracer graph and refresh decomposed outputs.

labels_to_ids_impl(module, analysis_batch, batch)

Implementation for converting string labels to tensor IDs.

logit_diffs_impl(module, analysis_batch, batch)

Implementation for computing logit differences.

model_ablation_impl(module, analysis_batch, ...)

Implementation for model ablation analysis.

model_cache_forward_impl(module, ...)

Implementation for forward pass with activation caching and latent model (SAE) hooks.

model_forward_impl(module, analysis_batch, ...)

Implementation for basic model forward pass.

model_fwd_impl(module, analysis_batch, ...)

Implementation for basic model forward pass.

model_fwd_intervention_impl(module, ...)

Apply generalized hook-point interventions and return pre/post logits.

model_fwd_w_cache_impl(module, ...)

Implementation for forward pass with activation caching (no latent model hooks).

model_fwd_w_cache_latent_models_impl(module, ...)

Implementation for forward pass with activation caching and latent model (SAE) hooks.

model_gradient_impl(module, analysis_batch, ...)

Implementation for gradient-based attribution.

sae_correct_acts_impl(module, ...)

Implementation for computing correct activations from SAE outputs.
interpretune.analysis.ops.definitions.ablate_sae_latent(sae_acts, hook, latent_idx=None, seq_pos=None)[source]#

Ablate a particular latent at a particular sequence position.

If either argument is None, we ablate at all latents / sequence positions.

Return type:

Tensor

Parameters:

  • sae_acts (Tensor)

  • hook (transformer_lens.hook_points.HookPoint)

  • latent_idx (int | None)

  • seq_pos (Tensor | None)

interpretune.analysis.ops.definitions.ablation_attribution_impl(module, analysis_batch, batch, logit_diff_fn=<function boolean_logits_to_avg_logit_diff>, get_loss_preds_diffs=<function get_loss_preds_diffs>)[source]#

Implementation for computing attribution values using latent ablation.

Return type:

DefaultAnalysisBatchProtocol

Parameters:
interpretune.analysis.ops.definitions.boolean_logits_to_avg_logit_diff(logits, target_indices, reduction=None)[source]#

Returns the avg logit diff on a set of prompts, with fixed s2 pos and stuff.

Return type:

Tensor

Parameters:

  • logits (Float[Tensor, 'batch seq 2'])

  • target_indices (Tensor)

  • reduction (Literal['mean', 'sum'] | None)

interpretune.analysis.ops.definitions.compute_attribution_graph_impl(module, analysis_batch, batch, batch_idx, **kwargs)[source]#

Generate and decompose a circuit-tracer attribution graph.

Return type:

AnalysisBatch

Parameters:
interpretune.analysis.ops.definitions.concept_direction_impl(module, analysis_batch, batch, batch_idx, **kwargs)[source]#

Compute a concept direction from latent-example rows, or fall back to token-group embeddings.

Aggregation modes (selected via concept_aggregate_output_mode on the batch):

  • "streaming": maintain per-group running weighted state sums and weight totals on analysis_inputs.store (concept_running_state_sum_a, concept_running_weight_a, concept_running_state_sum_b, concept_running_weight_b). Each per-batch invocation updates the running aggregates from this batch’s per-batch concept_latent_state / concept_group_id / concept_example_weight tensors and recomputes the current concept direction from the accumulated sums; the final batch’s emitted direction is the converged result. Memory cost is O(d_model * num_groups) instead of O(num_examples * d_model). For paired_rejection, additionally maintains pending per-group buffers (concept_pending_{a,b}_states, concept_pending_{a,b}_weights) plus running residual and pair-weight totals (concept_running_residual_sum, concept_running_pair_weight); pairs are matched by stable iteration order (matching the legacy in_memory contract). The full storage-contract field set is exported via interpretune.analysis.ops.helpers.CONCEPT_STREAMING_STATE_FIELDS.

  • "in_memory" (legacy): consume the aggregate row tensors emitted by extract_concept_latent_examples_impl() in legacy mode and compute the direction over the full materialized example set. Supports all direction modes.

If concept_aggregate_output_mode is not set, behavior is determined by what is on the batch: aggregate row fields trigger the legacy path; per-batch fields with a writable store trigger streaming. If neither is available, fall back to a token-group embedding direction computed from the model’s input embedding matrix.

Return type:

AnalysisBatch

Parameters:
interpretune.analysis.ops.definitions.extract_concept_latent_examples_impl(module, analysis_batch, batch, batch_idx, **kwargs)[source]#

Filter and annotate concept latent rows for downstream concept-direction aggregation.

Consumes concept_latent_state rows produced by the upstream extract_concept_latent_state op. That op must run first to populate concept_latent_state on the batch.

Aggregation modes (selected via concept_aggregate_output_mode on the batch):

  • "streaming" (default): emit only per-batch tensors. Cross-batch aggregation is performed incrementally inside concept_direction_impl() using running per-group weighted sums stored on analysis_inputs.store. This avoids materializing the full latent-row payload and keeps per-batch payload sizes constant. Recommended for all new callers and any pipeline where the full set of selected examples does not need to be retained for later inspection.

  • "in_memory" (legacy): accumulate the full per-batch row collections on analysis_inputs.store (concept_latent_state_rows, concept_group_id_rows, concept_group_name_rows, concept_example_logit_diff_rows, concept_example_weight_rows, optionally concept_context_indices_rows) and re-emit them on every returned batch. Each per-batch call appends to a Python list and re-binds it on the store and on analysis_batch; the underlying tensor data is shared by reference, but the list overhead grows linearly per batch and the runner’s per-batch payloads end up holding O(N²) cumulative list references for N batches. This mode remains useful for callers that need access to the full row collection (e.g. parity tests, pre-computed aggregate inputs to concept_direction), but its scalability depends on store caching/persistence; do not use it for large concept-example sets without that backing.

The legacy mode is preserved to keep existing tests and notebook diagnostics that consume the aggregate row tensors directly working unchanged.

Return type:

AnalysisBatch

Parameters:
interpretune.analysis.ops.definitions.extract_concept_latent_state_impl(module, analysis_batch, batch, batch_idx, **kwargs)[source]#

Extract per-example latent rows from the configured cache key for downstream concept-direction ops.

Return type:

AnalysisBatch

Parameters:
interpretune.analysis.ops.definitions.extract_top_features_impl(module, analysis_batch, batch, batch_idx, top_n=None, **kwargs)[source]#

Extract the top scoring features from analysis-batch feature rows.

An optional feature_selection kwarg (FeatureSelectionSpec) pre-filters active_features rows before score sorting. The filter uses OR semantics — a row is kept if it matches any criterion in the spec.

Return type:

AnalysisBatch

Parameters:
interpretune.analysis.ops.definitions.feature_intervention_forward_impl(module, analysis_batch, batch, batch_idx, **kwargs)[source]#

Run circuit-tracer feature interventions against the module replacement model.

This op currently implements forward-only intervention analysis and stores both the circuit-tracer tuple payload and a canonical feature-target InterventionDict summary for AnalysisStore consumers.

Return type:

AnalysisBatch

Parameters:
interpretune.analysis.ops.definitions.get_alive_latents_impl(module, analysis_batch, batch_idx)[source]#

Implementation for extracting alive latents from cache.

Return type:

DefaultAnalysisBatchProtocol

Parameters:
interpretune.analysis.ops.definitions.get_answer_indices_impl(module, analysis_batch, batch, batch_idx)[source]#

Implementation for extracting answer indices from batch.

Return type:

DefaultAnalysisBatchProtocol

Parameters:
interpretune.analysis.ops.definitions.get_loss_preds_diffs(module, analysis_batch, answer_logits, logit_diff_fn=<function boolean_logits_to_avg_logit_diff>)[source]#

Implementation for computing loss, predictions, and logit differences.

Parameters:

  • module (Module) – The module containing loss_fn and standardize_logits methods

  • analysis_batch (DefaultAnalysisBatchProtocol) – The analysis batch containing labels and orig_labels

  • answer_logits (Tensor) – The logits to analyze

  • logit_diff_fn (Callable) – Function to compute logit differences

Return type:

tuple[Tensor, Tensor, Tensor, Tensor]

Returns:

Tuple of (loss, logit_diffs, preds, answer_logits)

interpretune.analysis.ops.definitions.gradient_attribution_impl(module, analysis_batch, batch, batch_idx)[source]#

Implementation for computing attribution values from gradients.

Return type:

DefaultAnalysisBatchProtocol

Parameters:
interpretune.analysis.ops.definitions.graph_node_influence_impl(module, analysis_batch, batch, batch_idx, **kwargs)[source]#

Compute feature-node influence scores from a structured graph.

Return type:

AnalysisBatch

Parameters:
interpretune.analysis.ops.definitions.graph_prune_impl(module, analysis_batch, batch, batch_idx, **kwargs)[source]#

Prune a structured circuit-tracer graph and refresh decomposed outputs.

Return type:

AnalysisBatch

Parameters:
interpretune.analysis.ops.definitions.labels_to_ids_impl(module, analysis_batch, batch)[source]#

Implementation for converting string labels to tensor IDs.

Return type:

DefaultAnalysisBatchProtocol

Parameters:
interpretune.analysis.ops.definitions.logit_diffs_impl(module, analysis_batch, batch, logit_diff_fn=<function boolean_logits_to_avg_logit_diff>, get_loss_preds_diffs=<function get_loss_preds_diffs>)[source]#

Implementation for computing logit differences.

Return type:

DefaultAnalysisBatchProtocol

Parameters:
interpretune.analysis.ops.definitions.model_ablation_impl(module, analysis_batch, batch, batch_idx, ablate_latent_fn=<function ablate_sae_latent>)[source]#

Implementation for model ablation analysis.

Return type:

DefaultAnalysisBatchProtocol

Parameters:
interpretune.analysis.ops.definitions.model_cache_forward_impl(module, analysis_batch, batch, batch_idx)#

Implementation for forward pass with activation caching and latent model (SAE) hooks.

Return type:

DefaultAnalysisBatchProtocol

Parameters:
interpretune.analysis.ops.definitions.model_forward_impl(module, analysis_batch, batch, batch_idx)#

Implementation for basic model forward pass.

Return type:

DefaultAnalysisBatchProtocol

Parameters:
interpretune.analysis.ops.definitions.model_fwd_impl(module, analysis_batch, batch, batch_idx)[source]#

Implementation for basic model forward pass.

Return type:

DefaultAnalysisBatchProtocol

Parameters:
interpretune.analysis.ops.definitions.model_fwd_intervention_impl(module, analysis_batch, batch, batch_idx, **kwargs)[source]#

Apply generalized hook-point interventions and return pre/post logits.

Delegates the intervention mechanics to the model backend’s fwd_w_intervention method so the same op works for both NNsight (traced execution) and TransformerLens (eager hook execution).

Return type:

AnalysisBatch

Parameters:
interpretune.analysis.ops.definitions.model_fwd_w_cache_impl(module, analysis_batch, batch, batch_idx)[source]#

Implementation for forward pass with activation caching (no latent model hooks).

Return type:

DefaultAnalysisBatchProtocol

Parameters:
interpretune.analysis.ops.definitions.model_fwd_w_cache_latent_models_impl(module, analysis_batch, batch, batch_idx)[source]#

Implementation for forward pass with activation caching and latent model (SAE) hooks.

Return type:

DefaultAnalysisBatchProtocol

Parameters:
interpretune.analysis.ops.definitions.model_gradient_impl(module, analysis_batch, batch, batch_idx, logit_diff_fn=<function boolean_logits_to_avg_logit_diff>, get_loss_preds_diffs=<function get_loss_preds_diffs>)[source]#

Implementation for gradient-based attribution.

Defines a backward_fn closure that extracts answer logits, computes logit diffs, and returns their sum as the scalar to backpropagate. The backend handles the entire forward + backward flow (enabling both eager and trace-based execution).

Return type:

DefaultAnalysisBatchProtocol

Parameters:
interpretune.analysis.ops.definitions.sae_correct_acts_impl(module, analysis_batch, batch, batch_idx)[source]#

Implementation for computing correct activations from SAE outputs.

Return type:

DefaultAnalysisBatchProtocol

Parameters: