from __future__ import annotations # see PEP 749, no longer needed when 3.13 reaches EOL
from typing import Any, Generator, Callable
from dataclasses import dataclass, field
import datetime
import warnings
from weakref import WeakKeyDictionary
from transformers import BatchEncoding, PreTrainedTokenizerBase
from interpretune.analysis import (
LatentAnalysisTargets,
resolve_names_filter,
_make_simple_cache_hook,
OpSchema,
AnalysisOp,
OpWrapper,
AnalysisOpLike,
)
from interpretune.analysis.execution import execute_analysis_step
from interpretune.analysis.ops.dispatcher import DISPATCHER
from interpretune.config import ITSerializableCfg
from interpretune.protocol import NamesFilter, AnalysisStoreProtocol, BaseAnalysisBatchProtocol, STEP_OUTPUT
from interpretune.utils import DEFAULT_DECODE_KWARGS
from interpretune.utils import rank_zero_warn, rank_zero_debug
def _extend_names_for_bridge(module, names_list: list[str]) -> tuple[list[str], dict[str, str]]:
"""Extend a names_filter list with canonical-name equivalents for TransformerBridge models.
:func:`construct_names_filter` builds alias-based hook names (e.g.
``blocks.9.attn.hook_z.hook_sae_acts_post``) but TransformerBridge's
``hook_dict`` and ``run_with_cache`` use canonical names (e.g.
``blocks.9.attn.o.hook_in.hook_sae_acts_post``). When the resolved callable
is passed to ``run_with_cache``, Bridge doesn't resolve aliases for callables
(only for list inputs). This helper adds the canonical equivalents to the
filter list so the callable matches both alias and canonical keys.
For non-Bridge models this is a no-op that returns the input list unchanged
and an empty mapping.
Args:
module: The IT module whose ``model`` attribute is checked for Bridge type.
names_list: Original names_filter list with alias-based names.
Returns:
Tuple of (extended_list, canonical_to_alias_map) where:
- extended_list includes canonical-name equivalents (or original list for non-Bridge)
- canonical_to_alias_map maps canonical names back to their alias originals
(empty dict for non-Bridge models)
"""
try:
from transformer_lens.model_bridge.bridge import TransformerBridge
except ImportError:
return names_list, {}
model = getattr(module, "model", None)
if model is None or not isinstance(model, TransformerBridge):
return names_list, {}
# Build combined alias→canonical mapping from the bridge.
# ``hook_aliases`` contains static aliases (embed/pos_embed/unembed).
# ``_collect_hook_aliases_from_registry`` (on SAETransformerBridge) contains
# dynamic aliases including the critical hook_z→o.hook_in mappings.
alias_to_canonical: dict[str, str] = {}
if hasattr(model, "hook_aliases"):
alias_to_canonical.update(model.hook_aliases) # type: ignore[arg-type] # TL hook_aliases is Dict[str, str | List[str]]
if hasattr(model, "_collect_hook_aliases_from_registry"):
alias_to_canonical.update(model._collect_hook_aliases_from_registry()) # type: ignore[arg-type]
if not alias_to_canonical:
return names_list, {}
extended = list(names_list)
canonical_to_alias: dict[str, str] = {}
for name in names_list:
# Try each alias prefix: if the name starts with an alias prefix (possibly with a suffix),
# add the canonical equivalent. e.g. "blocks.9.attn.hook_z.hook_sae_acts_post"
# → alias "blocks.9.attn.hook_z" maps to "blocks.9.attn.o.hook_in"
# → canonical equiv: "blocks.9.attn.o.hook_in.hook_sae_acts_post"
for alias, canonical in alias_to_canonical.items():
if name == alias:
if canonical not in extended:
extended.append(canonical)
canonical_to_alias[canonical] = name
break
elif name.startswith(alias + "."):
suffix = name[len(alias) :]
canonical_name = canonical + suffix
if canonical_name not in extended:
extended.append(canonical_name)
canonical_to_alias[canonical_name] = name
break
return extended, canonical_to_alias
[docs]
@dataclass(kw_only=True)
class AnalysisCfg(ITSerializableCfg):
output_store: AnalysisStoreProtocol | None = None # usually constructed on setup()
input_store: AnalysisStoreProtocol | None = None # store containing input data from previous op
batch_inputs: dict[str, Any] = field(default_factory=dict) # batch-scoped inputs for the active invocation
run_inputs: dict[str, Any] = field(default_factory=dict) # run-scoped inputs shared across the invocation
target_op: str | AnalysisOp | Callable | list[AnalysisOp | Callable] | None = None # input op to be resolved
output_schema: OpSchema | str | AnalysisOp | Callable | None = None # Schema, op, or op name to define schema
name: str | None = None # Name for this analysis configuration
fwd_hooks: list[tuple] = field(default_factory=list)
bwd_hooks: list[tuple] = field(default_factory=list)
cache_dict: dict = field(default_factory=dict)
names_filter: NamesFilter | None = None
save_prompts: bool = False
save_tokens: bool = False
decode_kwargs: dict = field(default_factory=lambda: DEFAULT_DECODE_KWARGS)
latent_analysis_targets: LatentAnalysisTargets | None = None
ignore_manual: bool = False # When True, ignore existing analysis_step and use op to generate one
step_fn: str = "analysis_step" # Name of the method to use/generate for analysis
# Preserved original step_fn name for idempotent re-application across module instances
# (prevents _generated_ prefix accumulation when this cfg is shared via class-level defaults)
_original_step_fn: str | None = field(default=None, init=False, repr=False, compare=False)
auto_prune_batch_encoding: bool = True # Automatically prune encoded batches to only include relevant keys
_applied_to: WeakKeyDictionary = field( # type: ignore[type-arg]
default_factory=WeakKeyDictionary,
init=False,
repr=False,
compare=False,
) # Tracks which live module instances this cfg has been applied to
_op: str | AnalysisOp | Callable | list[AnalysisOp | Callable] | None = None # op via generated analysis step
@property
def op(self) -> str | AnalysisOp | Callable | list[AnalysisOp | Callable] | None:
"""Get the operation, unwrapping any OpWrapper if present."""
if self._op is None:
return None
# Unwrap OpWrapper instances if needed - only check on objects that might have these attributes
if isinstance(self._op, OpWrapper) and getattr(self._op, "_is_instantiated", False):
return getattr(self._op, "_instantiated_op")
return self._op
@op.setter
def op(self, value: str | AnalysisOp | Callable | list[AnalysisOp | Callable] | None) -> None:
"""Set the operation value."""
self._op = value
def __post_init__(self):
# Process the target_op if provided and set it as the op
if self.target_op is not None:
# TODO: consider saving the original target_op value before assigning to self.op (str, op, wrapper, etc.)
# and then after the op is resolved by the end of the post_init, replace self.target_op with a
# well-formatted string representation of the original target_op value and what it was resolved to
# (e.g. "target_op: 'some_op_str' -> op: {str(self.op)}")
self.op = self.target_op
if self.ignore_manual and self.op is None: # Check if ignore_manual is True but no op is provided
raise ValueError("When ignore_manual is True, an op must be provided to generate the analysis_step")
if self.output_schema is not None and not isinstance(self.output_schema, OpSchema):
self.resolve_output_schema() # Resolve the output schema if it's not already an OpSchema
if self.op is not None: # Process the operation if provided
self.resolve_op()
if self.name is None and self.op is not None: # Set name from op if name is not already set
if hasattr(self.op, "name"):
self.name = getattr(self.op, "name")
if self.name is None: # If name still not set (no op was resolved), use timestamp default
self.name = f"default_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}"
[docs]
def update(self, **kwargs):
"""Update multiple fields of the dataclass at once."""
for key, value in kwargs.items():
if key in self.__annotations__:
setattr(self, key, value)
def resolve_output_schema(self) -> None:
assert self.output_schema is not None, "Output schema to resolve must be set"
# Accept either AnalysisOp or OpWrapper-like objects that will resolve to AnalysisOp when instantiated
if isinstance(self.output_schema, AnalysisOpLike):
resolved_op = self.output_schema
self.output_schema = resolved_op.output_schema
# If output_schema is a string, resolve to op and extract schema
elif isinstance(self.output_schema, str):
resolved_op = DISPATCHER.get_op(self.output_schema)
assert hasattr(resolved_op, "output_schema"), "Resolved op is malformed"
self.output_schema = getattr(resolved_op, "output_schema")
def resolve_op(self) -> None:
# Handle list inputs (composition)
# TODO: this resolution still needs to be both refactored (we haven't fully refined excessively defensive
# AI written conditions) and enhanced to handle additional not officially supported but possible to support
# cases (e.g. OpWrapper instance in a list,)
if isinstance(self.op, list): # Convert list of ops to composition
if len(self.op) == 1:
self.op = self.op[0]
if isinstance(self.op, str):
self.op = DISPATCHER.get_op(self.op)
else:
# Create a composite operation, ensuring constituent ops are instantiated
# TODO: consider deferring instantiation of composite ops
instantiated_ops = []
for op in self.op:
if isinstance(op, str):
instantiated_ops.append(DISPATCHER.get_op(op))
elif isinstance(op, OpWrapper):
instantiated_ops.append(op._ensure_instantiated())
elif isinstance(op, AnalysisOp):
instantiated_ops.append(op)
else:
ensure_instantiated = getattr(op, "_ensure_instantiated", None)
if callable(ensure_instantiated):
instantiated_ops.append(ensure_instantiated())
continue
bound_self = getattr(op, "__self__", None)
bound_ensure_instantiated = getattr(bound_self, "_ensure_instantiated", None)
if callable(bound_ensure_instantiated):
instantiated_ops.append(bound_ensure_instantiated())
continue
op_name = (
getattr(op, "_op_name", None) or getattr(op, "name", None) or getattr(op, "__name__", None)
)
if isinstance(op_name, str):
try:
instantiated_ops.append(DISPATCHER.get_op(op_name))
continue
except Exception:
pass
instantiated = op()
if isinstance(instantiated, OpWrapper):
instantiated_ops.append(instantiated._ensure_instantiated())
elif isinstance(instantiated, AnalysisOp):
instantiated_ops.append(instantiated)
else:
raise TypeError(f"Composite op factory returned unsupported type: {type(instantiated)}")
self.op = DISPATCHER.compile_ops(instantiated_ops)
return
# Handle string names or composite names using dot notation
if isinstance(self.op, str):
if "." in self.op:
# This is a composite operation name
self.op = DISPATCHER.compile_ops(self.op)
else:
# Resolve single op name or alias
self.op = DISPATCHER.get_op(self.op)
return
# Defensive normalization for non-instance inputs (classes, wrappers, etc.)
candidate = self._op
# Try common attributes for an op name
op_name = getattr(candidate, "_op_name", None) or getattr(candidate, "name", None)
# If a class/type was passed, try to resolve an op name or instantiate
if isinstance(candidate, type):
if isinstance(op_name, str):
try:
self.op = DISPATCHER.get_op(op_name)
return
except Exception:
# fall through to other heuristics
pass
# Try to call the class (if it returns a wrapper or op instance)
try:
inst = candidate()
except Exception:
inst = None
if isinstance(inst, (OpWrapper, AnalysisOp)):
# Replace stored value with the instantiated object
self._op = inst
if isinstance(inst, OpWrapper):
try:
resolved = inst._ensure_instantiated()
self._op = resolved
except Exception:
# keep wrapper instance if unable to fully resolve
self._op = inst
return
# If provided object has an _op_name or name attribute, try resolving via dispatcher
if isinstance(op_name, str):
try:
self.op = DISPATCHER.get_op(op_name)
return
except Exception:
pass
# Otherwise leave as-is; callers will raise clear errors if this is invalid
[docs]
def materialize_names_filter(self, module, fallback_sae_targets: LatentAnalysisTargets | None = None) -> None:
"""Set names_filter using latent_analysis_targets if not already set.
For :class:`~transformer_lens.model_bridge.bridge.TransformerBridge` models the constructed
filter is automatically extended with canonical-name equivalents so that both
``run_with_cache`` (which iterates canonical ``hook_dict`` keys) and downstream filtering
(e.g. ``get_alive_latents_impl``) can match activations regardless of whether the cache
uses alias or canonical naming.
Args:
module: The module to construct the names_filter for.
fallback_sae_targets: Optional fallback LatentAnalysisTargets to use if this config doesn't have one.
"""
# Skip if names_filter is already set
if self.names_filter is not None:
self.names_filter = resolve_names_filter(self.names_filter)
return
# Choose the appropriate LatentAnalysisTargets
sae_targets = self.latent_analysis_targets or fallback_sae_targets
if sae_targets is not None:
target_layers = sae_targets.target_layers
match_fn = sae_targets.sae_hook_match_fn
self.names_filter = module.construct_names_filter(target_layers, match_fn)
else:
raise ValueError("No LatentAnalysisTargets available to create names_filter")
# For TransformerBridge models, extend the filter list with canonical-name equivalents.
# construct_names_filter builds alias-based names (e.g. blocks.9.attn.hook_z.hook_sae_acts_post)
# but Bridge's hook_dict and run_with_cache use canonical names
# (e.g. blocks.9.attn.o.hook_in.hook_sae_acts_post).
# Without this, the callable names_filter doesn't match canonical names → cache is empty →
# alive_latents empty.
# TODO: revisit names_filter handling for TransformerBridge — currently we extend the list
# with both alias and canonical names and remap keys post-hoc (_remap_bridge_hook_keys). A
# cleaner approach would be to resolve the naming scheme once at config time so downstream
# code only ever sees one consistent set of hook names.
if isinstance(self.names_filter, list):
self.names_filter, self._canonical_to_alias_names = _extend_names_for_bridge(module, self.names_filter)
else:
self._canonical_to_alias_names = {}
self.names_filter = resolve_names_filter(self.names_filter)
[docs]
def maybe_set_hooks(self) -> None:
"""Set hooks if they're not already set."""
if not self.fwd_hooks and not self.bwd_hooks:
self.check_add_default_hooks()
[docs]
def prepare_model_ctx(self, module, fallback_sae_targets: LatentAnalysisTargets | None = None) -> None:
"""Configure names_filter and hooks for a specific module.
Args:
module: The module to configure for.
fallback_sae_targets: Optional fallback LatentAnalysisTargets to use if this config doesn't have one.
"""
# Ensure latent_analysis_targets, fallback_sae_targets, or names_filter are set before materializing
if not (self.latent_analysis_targets or fallback_sae_targets or self.names_filter):
rank_zero_debug(
"None of (latent_analysis_targets, fallback_sae_targets, names_filter) are set. "
"Proceeding without materializing names_filter."
)
else:
# Always materialize names_filter - needed for both op-based and manual analysis
self.materialize_names_filter(module, fallback_sae_targets)
# Set hooks based on op or manually if needed
if self.op is not None:
self.maybe_set_hooks()
# TODO: Add a non-generator returning save_batch method at AnalysisCfg level (users already have op-level version)?
[docs]
def save_batch(
self,
analysis_batch: BaseAnalysisBatchProtocol,
batch: BatchEncoding,
tokenizer: PreTrainedTokenizerBase | None = None,
):
"""Process and yield analysis batch results.
Uses AnalysisOp.process_batch for consistent processing regardless of
whether an operation is defined or using a manual analysis step.
Args:
analysis_batch: The analysis batch to process
batch: The raw batch data
tokenizer: Optional tokenizer for decoding prompts
Yields:
Processed analysis batch
"""
if self.op is not None:
# When using a defined operation
assert isinstance(self.op, AnalysisOp), "op expected to be an instance of AnalysisOp"
analysis_batch = self.op.save_batch(
analysis_batch,
batch,
tokenizer=tokenizer,
save_prompts=self.save_prompts,
save_tokens=self.save_tokens,
decode_kwargs=self.decode_kwargs,
)
else:
# For manual analysis steps, use the static method with output_schema
from interpretune.analysis.ops.base import AnalysisOpLike
# Convert OpWrapper or AnalysisOp to OpSchema if needed
output_schema = self.output_schema
if isinstance(output_schema, AnalysisOpLike):
output_schema = output_schema.output_schema
assert isinstance(output_schema, (OpSchema, type(None))), "if set output_schema should be an OpSchema"
analysis_batch = AnalysisOp.process_batch(
analysis_batch,
batch,
output_schema=output_schema or OpSchema({}), # Ensure type checker knows this is OpSchema
tokenizer=tokenizer,
save_prompts=self.save_prompts,
save_tokens=self.save_tokens,
decode_kwargs=self.decode_kwargs,
)
# For TransformerBridge models, remap canonical hook-name keys to alias-based
# keys at the serialization boundary so the HF dataset matches its schema.
analysis_batch = self._remap_bridge_hook_keys(analysis_batch)
yield analysis_batch
def _remap_bridge_hook_keys(self, analysis_batch: BaseAnalysisBatchProtocol) -> BaseAnalysisBatchProtocol:
"""Remap canonical hook-name dict keys to alias-based keys for Bridge models.
TransformerBridge models use canonical hook names internally (e.g.
``blocks.9.attn.o.hook_in.hook_sae_acts_post``) while the dataset schema
uses alias-based names from SAE metadata (e.g.
``blocks.9.attn.hook_z.hook_sae_acts_post``).
This remapping is applied once at the serialization boundary so all
internal analysis operations can use consistent canonical names from the
activation cache.
Args:
analysis_batch: The analysis batch whose per-hook dict fields may need
key remapping.
Returns:
The analysis batch with remapped keys (modified in-place via setattr).
"""
canonical_to_alias = getattr(self, "_canonical_to_alias_names", None)
if not canonical_to_alias:
return analysis_batch
# Determine output schema to identify per-hook dict fields
from interpretune.analysis.ops.base import AnalysisOpLike, OpSchema
schema = None
if self.op is not None and hasattr(self.op, "output_schema"):
schema = self.op.output_schema # type: ignore[union-attr] # guarded by hasattr
elif hasattr(self, "output_schema") and self.output_schema is not None:
schema = self.output_schema
if isinstance(schema, AnalysisOpLike):
schema = schema.output_schema
if not isinstance(schema, OpSchema):
return analysis_batch
for col_name, col_cfg in schema.items():
if col_cfg.intermediate_only:
continue
if not (col_cfg.per_latent_model_hook or col_cfg.per_latent):
continue
val = getattr(analysis_batch, col_name, None)
if not isinstance(val, dict):
continue
remapped = {canonical_to_alias.get(k, k): v for k, v in val.items()}
setattr(analysis_batch, col_name, remapped)
return analysis_batch
[docs]
def add_default_cache_hooks(self, include_backward: bool = True) -> None:
"""Add default caching hooks for forward and optionally backward passes.
Args:
names_filter: Filter to determine which layers to hook
cache_dict: Dictionary to store activation values
include_backward: Whether to include backward hooks
Returns:
Tuple of (forward hooks, backward hooks)
"""
fwd_hooks = [(self.names_filter, _make_simple_cache_hook(cache_dict=self.cache_dict))]
self.fwd_hooks = fwd_hooks
bwd_hooks = []
if include_backward:
bwd_hooks = [(self.names_filter, _make_simple_cache_hook(cache_dict=self.cache_dict, is_backward=True))]
self.bwd_hooks = bwd_hooks
[docs]
def check_add_default_hooks(self) -> tuple[list, list] | None:
"""Construct forward and backward hooks based on analysis operation."""
fwd_hooks, bwd_hooks = [], []
if self.op is None:
self.fwd_hooks, self.bwd_hooks = fwd_hooks, bwd_hooks
return
assert isinstance(self.op, AnalysisOpLike), (
f"op expected to be AnalysisOp (or uninstantiated OpWrapper), got {type(self.op)}"
)
# TODO: change these op-based checks to be functionally driven (e.g. uses_default_hooks attribute of ops)
if self.op.name == "logit_diffs_base":
return fwd_hooks, bwd_hooks
if self.op.name == "logit_diffs_attr_grad":
self.add_default_cache_hooks()
# TODO: add in an op attribute akin to "uses_sae_hooks" to enable names_filter validation resolution etc.
# if self.op_name in ('logit_diffs_attr_ablation', 'logit_diffs_attr_grad') and self.names_filter is None:
# raise ValueError("names_filter required for non-clean operations")
[docs]
def applied_to(self, module) -> bool:
"""Check if this configuration has been applied to a specific module.
Args:
module: The module to check.
Returns:
bool: True if this configuration has been applied to the module, False otherwise.
"""
return module in self._applied_to
[docs]
def reset_applied_state(self, module=None) -> None:
"""Reset the applied state tracking.
Args:
module: Optional specific module to reset. If None, reset for all modules.
"""
if module is not None:
# Reset for specific module
self._applied_to.pop(module, None)
else:
# Reset for all modules
self._applied_to.clear()
[docs]
def apply(
self,
module,
cache_dir: str | None = None,
op_output_dataset_path: str | None = None,
fallback_sae_targets: LatentAnalysisTargets | None = None,
):
"""Set up analysis configuration and configure for the given module.
This method handles both setting up the analysis store and configuring
the names_filter and hooks for the module. It also injects an analysis_step
method if one doesn't exist in the module.
Args:
module: The module to configure for.
cache_dir: Optional cache directory.
op_output_dataset_path: Optional output path.
fallback_sae_targets: Optional fallback LatentAnalysisTargets to use if this config doesn't have one.
"""
# Short-circuit if already applied to this module (though apply should be idempotent)
if module in self._applied_to:
return
# Store original step_fn name before any mutation to prevent _generated_ prefix accumulation
# when this AnalysisCfg instance is shared across multiple module instantiations
if self._original_step_fn is None:
self._original_step_fn = self.step_fn
base_step_fn = self._original_step_fn
# Check if module has a manually-defined analysis step (using the original/base name)
has_custom_step = hasattr(module, base_step_fn) and not getattr(module, f"_generated_{base_step_fn}", False)
# Only warn about custom step if we're not going to ignore it
if has_custom_step and not self.ignore_manual:
warnings.warn(
f"Module {module.__class__.__name__} already has a {base_step_fn} method. "
"The provided operation configuration will be used for hooks and filters, "
f"but the execution flow will be determined by the existing {base_step_fn} method."
)
# Generate a new step if no custom step exists OR we're explicitly ignoring manual steps
if (not has_custom_step or self.ignore_manual) and self.op is not None:
# Inject a dynamic analysis_step method
def generated_analysis_step(
self, batch: BatchEncoding, batch_idx: int, dataloader_idx: int = 0
) -> Generator[STEP_OUTPUT, None, None]:
"""Dynamically generated analysis_step method."""
yield from execute_analysis_step(
self,
batch,
batch_idx,
dataloader_idx=dataloader_idx,
analysis_cfg=self.analysis_cfg,
)
# TODO: separate some of this more ephemeral state to an AnalysisState object
# Add the method to the module with a _generated version of the base step_fn name
# (to avoid potentially clobbering the manual version; always uses base_step_fn so
# the generated name is stable even when this cfg is shared across module instances)
setattr(module, f"_generated_{base_step_fn}", generated_analysis_step.__get__(module))
# update the analysis_cfg step_fn to the generated name (idempotent: always single prefix)
setattr(self, "step_fn", f"_generated_{base_step_fn}")
# Always set up the analysis store, even for manual analysis steps
if not self.output_store:
# Create the output store if needed
from interpretune.analysis import AnalysisStore
self.output_store = AnalysisStore(cache_dir=cache_dir, op_output_dataset_path=op_output_dataset_path)
elif cache_dir or op_output_dataset_path:
rank_zero_warn(
f"The provided cache_dir={cache_dir} and op_output_dataset_path={op_output_dataset_path} "
"will be ignored in favor of the existing AnalysisStore configuration "
f"(cache_dir={self.output_store.cache_dir}, "
f" op_output_dataset_path={self.output_store.op_output_dataset_path})"
)
# Always prepare the model context to ensure names_filter is materialized
self.prepare_model_ctx(module, fallback_sae_targets)
# Mark as applied to this specific module, storing module class name for debugging
self._applied_to[module] = module.__class__.__name__
[docs]
@dataclass(kw_only=True)
class AnalysisArtifactCfg(ITSerializableCfg):
"""Configuration for analysis artifacts and visualizations."""
latent_effects_graphs: bool = True
latent_effects_graphs_per_batch: bool = False # can be overwhelming with many batches
table_per_latent_model: bool = True
top_k_latents_table: int = 2
top_k_latent_dashboards: int = 1 # (don't set too high, num dashboards = top_k_latent_dashboards * num_hooks * 2)
top_k_clean_logit_diffs: int = 10
def __post_init__(self):
if self.latent_effects_graphs_per_batch and not self.latent_effects_graphs:
print("Note: Setting latent_effects_graphs to True since latent_effects_graphs_per_batch is True")
self.latent_effects_graphs = True
