interpretune.analysis.ops.base#

Base classes for analysis operations.

Functions

`build_call_args`(module, analysis_batch, ...)	Build arguments for operation calls.
`capability_value`(capability)	Return a stable string key for a backend capability-like object.
`get_batch_input`(batch)	Get the tokenized input tensor from the batch, handling cross-backend key differences.
`resolve_batch_input_key`(batch[, key])	Resolve the actual batch key, checking aliases for cross-backend compatibility.
`wrap_summary`(analysis_batch, batch[, ...])

Classes

`AnalysisBatch`(args, *kwargs)
`AnalysisOp`(name, description, output_schema)	Base class for analysis operations.
`AttrDict`(args, *kwargs)
`ColCfg`(datasets_dtype[, required, dyn_dim, ...])	Configuration for a dataset column.
`CompositeAnalysisOp`(ops[, name, aliases, ...])	A composition of analysis operations to be executed.
`OpSchema`(args, *kwargs)	Schema defining column specifications for analysis operations.
`OpWrapper`(op_name)	A special wrapper for operations that ensures the op is instantiated when accessed directly or when attributes are accessed.

class interpretune.analysis.ops.base.AnalysisBatch(*args, **kwargs)[source]#

bind_resolution_context(module, *, analysis_inputs=None, batch_idx=None, input_schema=None)[source]#

Bind execution-time lookup context used by scoped batch accessors.

Return type:

AnalysisBatch

Parameters:

module (Any | None)
analysis_inputs (Any)
batch_idx (int | None)
input_schema (Any)

clear_resolution_context()[source]#

Clear any previously bound execution-time lookup context.

Return type:: AnalysisBatch

get(key, default=None, *, scopes=None, resolve=True)[source]#

Return a direct batch value or, when bound, resolve via scoped inputs.

Parameters:

scopes (tuple[str, ...] | None)
resolve (bool)

require(key, *, scopes=None, message=None)[source]#

Resolve a required value from the bound analysis execution context.

Return type:

Any

Parameters:

key (str)
scopes (tuple[str, ...] | None)
message (str | None)

resolution_context(module, *, analysis_inputs=None, batch_idx=None, input_schema=None)[source]#

Temporarily bind execution-time scoped lookup state for the current op call.

Parameters:

module (Any | None)
analysis_inputs (Any)
batch_idx (int | None)
input_schema (Any)

resolve(key, default=None, *, scopes=None)[source]#

Resolve a value using the bound analysis execution context.

Return type:

Any

Parameters:

key (str)
default (Any)
scopes (tuple[str, ...] | None)

to_cpu()[source]#: Detach and move all field tensors to CPU.

update([E, ]**F) → None. Update D from mapping/iterable E and F.[source]#: If E is present and has a .keys() method, then does: for k in E.keys(): D[k] = E[k] If E is present and lacks a .keys() method, then does: for k, v in E: D[k] = v In either case, this is followed by: for k in F: D[k] = F[k]

class interpretune.analysis.ops.base.AnalysisOp(name, description, output_schema, input_schema=None, aliases=None, impl_params=None, required_capabilities=None)[source]#

Base class for analysis operations.

Parameters:

name (str)
description (str)
output_schema (OpSchema)
input_schema (OpSchema | None)
aliases (Sequence[str] | None)
impl_params (dict[str, Any] | None)
required_capabilities (Sequence[str | Any] | None)

static process_batch(analysis_batch, batch, output_schema, tokenizer=None, save_prompts=False, save_tokens=False, decode_kwargs=None)[source]#

Process analysis batch using provided output schema.

This static method handles the common processing logic for analysis batches, including token handling and schema-based transformations.

Parameters:

analysis_batch (BaseAnalysisBatchProtocol) – The analysis batch to process
batch (BatchEncoding) – The raw batch data
output_schema (OpSchema) – Schema defining the structure of the output
tokenizer (PreTrainedTokenizerBase | None) – Optional tokenizer for decoding prompts
save_prompts (bool) – Whether to save prompts
save_tokens (bool) – Whether to save tokens
decode_kwargs (dict[str, Any] | None) – Additional keyword arguments for decoding

Return type:

BaseAnalysisBatchProtocol

Returns:

Processed analysis batch

active_ctx_key(ctx_key)[source]#

Context manager for temporarily setting the active context key.

Parameters:: ctx_key – The context key to set during the context execution

save_batch(analysis_batch, batch, tokenizer=None, save_prompts=False, save_tokens=False, decode_kwargs=None)[source]#

Save analysis batch using process_batch static method.

Return type:

BaseAnalysisBatchProtocol

Parameters:

analysis_batch (BaseAnalysisBatchProtocol)
batch (BatchEncoding)
tokenizer (PreTrainedTokenizerBase | None)
save_prompts (bool)
save_tokens (bool)
decode_kwargs (dict[str, Any] | None)

property ctx_key: str#: Return the context key if set, otherwise return the name.

property impl: Callable | None#: Get the implementation function.

class interpretune.analysis.ops.base.AttrDict(*args, **kwargs)[source]#

class interpretune.analysis.ops.base.ColCfg(datasets_dtype, required=True, dyn_dim=None, dyn_dim_ceil=None, non_tensor=False, per_latent=False, per_latent_model_hook=False, intermediate_only=False, connected_obj='analysis_store', array_shape=None, sequence_type=True, array_dtype=None, default=None)[source]#

Configuration for a dataset column.

Parameters:

datasets_dtype (str)
required (bool)
dyn_dim (int | None)
dyn_dim_ceil (Literal['batch_size', 'max_answer_tokens', 'num_classes', 'vocab_size', 'max_seq_len'] | None)
non_tensor (bool)
per_latent (bool)
per_latent_model_hook (bool)
intermediate_only (bool)
connected_obj (Literal['analysis_store', 'datamodule'])
array_shape (tuple[int | Literal['batch_size', 'max_answer_tokens', 'num_classes', 'vocab_size', 'max_seq_len'] | None, ...] | None)
sequence_type (bool)
array_dtype (str | None)
default (Any)

classmethod from_dict(data)[source]#

Create from dict representation.

Return type:: ColCfg
Parameters:: data (dict)

to_dict()[source]#

Convert to JSON serializable dict.

Return type:: dict

class interpretune.analysis.ops.base.CompositeAnalysisOp(ops, name=None, aliases=None, description=None, *args, **kwargs)[source]#

A composition of analysis operations to be executed.

Parameters:

ops (Sequence[AnalysisOp])
name (str | None)
aliases (Sequence[str] | None)
description (str | None)

class interpretune.analysis.ops.base.OpSchema(*args, **kwargs)[source]#: Schema defining column specifications for analysis operations.

class interpretune.analysis.ops.base.OpWrapper(op_name)[source]#

A special wrapper for operations that ensures the op is instantiated when accessed directly or when attributes are accessed.

classmethod initialize(target_module)[source]#: Set the target module where operations will be registered.

classmethod register_operations(module, dispatcher)[source]#

Parameters:

module – The module where operations will be registered
dispatcher – The operations dispatcher instance

property dispatcher#: Lazily load the dispatcher only when needed.

interpretune.analysis.ops.base.build_call_args(module, analysis_batch, batch, batch_idx, impl_params=None, **kwargs)[source]#

Build arguments for operation calls.

Parameters:

module – The module instance
analysis_batch – The analysis batch
batch – The input batch
batch_idx – The batch index
impl_params – Implementation-specific parameters
**kwargs – Additional keyword arguments

Returns:

Dictionary of arguments for the operation call

interpretune.analysis.ops.base.capability_value(capability)[source]#

Return a stable string key for a backend capability-like object.

Return type:: str
Parameters:: capability (Any)

interpretune.analysis.ops.base.get_batch_input(batch)[source]#

Get the tokenized input tensor from the batch, handling cross-backend key differences.

Supports both TransformerLens convention (batch["input"]) and HuggingFace/NNsight convention (batch["input_ids"]).

Return type:: Tensor
Parameters:: batch (BatchEncoding)

interpretune.analysis.ops.base.resolve_batch_input_key(batch, key='input')[source]#

Resolve the actual batch key, checking aliases for cross-backend compatibility.

Parameters:

batch (BatchEncoding) – The batch to check for the key.
key (str) – The canonical key name to resolve (default: "input").

Return type:

str

Returns:

The actual key present in the batch (either the canonical key or one of its aliases).

Raises:

KeyError – If neither the key nor any of its aliases are present in the batch.