entoscope.circuit

Circuit discovery.

The operator, its input, configs, collectors, method base and factory are importable from entoscope.circuit; the operator, its input and configs from entoscope too.

class entoscope.circuit.config.CircuitMethodConfig(*, type: str)

Bases: MethodConfig

The method config of a circuit discovery method that takes no fields.

A method with fields of its own names a dataclass extending this one.

classmethod factory() → type[CircuitMethodFactory]

Name the factory circuit discovery methods register with.

Returns:

CircuitMethodFactory.

class entoscope.circuit.config.CircuitConfig(*, config_id: str | None = None, method: CircuitMethodConfig)

Bases: AbstractConfig

What a circuit discovery run builds.

Unlike the per-location operations (lens, patching and probing), circuit discovery is global: a method builds a graph spanning every layer, token and head, so there is no locator here. Whether a page is drawn is chosen by the collector, not here.

Variables:

method (entoscope.circuit.config.CircuitMethodConfig) – The method config, its type naming the method. In a YAML file it is the mapping read strictly into the method’s config class, or just the method’s type for its defaults.

class entoscope.circuit.input.CircuitInput(config: CircuitConfig, prompts: list[str])

Bases: OperationInput

The prompts to discover circuits for, one contribution graph per prompt.

It holds no activations: the method runs each prompt as it builds that prompt’s circuit, so peak memory stays bounded by the largest single prompt regardless of how many prompts are passed.

Variables:

• config (entoscope.circuit.config.CircuitConfig) – The circuit discovery config.

• prompts (list[str]) – The prompts.

Raises:

ValueError – If prompts is empty.

prepare(model: TransformerBridge | CachedModel, *, index_only: bool = False) → CircuitInput

Resolve the index a run iterates, then gather the activations it reads on top of it.

The index needs no forward pass, so index_only lets an operation be planned against a model without spending the compute. An input that gathers no activations, its method running the model itself, is ready once its index is.

Parameters:

• model – The model the run reads.

• index_only – Whether to stop after the index, before any forward pass.

Returns:

The input itself, prepared in place.

to(device: str | device = 'cpu') → CircuitInput

Move the tensors the input holds to device.

Parameters:

device – The device to move to.

Returns:

The input itself, moved in place.

class entoscope.circuit.discoverer.CircuitDiscoverer(model: TransformerBridge | CachedModel, collector_class: type[CircuitCollector] = CircuitVisualizationCollector)

Bases: Operator

Discovers the circuit behind each prompt with the method config.method names.

There is no index: a circuit spans every layer, position and head, so the method builds one graph per prompt from a forward pass of that prompt.

Parameters:

• model – The model to read.

• collector_class – The collector to write the results with. The default also draws a circuit_{i}.html per prompt, and CircuitCollector writes the graph files alone.

apply(circuit_input: CircuitInput, save_dir: str | None = None, return_collector: bool = False, verbose: bool = True) → CircuitCollector | None

Build the circuit behind every prompt of the input.

The method is built before any forward pass, so a config that names something the run cannot do fails before the model is touched. Where the results were saved is logged through the entoscope logger.

Parameters:

• circuit_input – The config and prompts to run.

• save_dir – The directory to write every prompt’s graph files and page, prompts.json and config.json into, created if missing. None writes nothing.

• return_collector – Whether to return the collector.

• verbose – Whether to show a progress bar.

Returns:

The collector, whose result holds one PromptCircuit per prompt, if return_collector is set, and None otherwise.

plan(circuit_input: CircuitInput, save_dir: str | None = None, verbose: bool = True) → OperationPlan

Describe the run apply would make, without a forward pass and without touching save_dir.

Parameters:

• circuit_input – The config and prompts the run would use.

• save_dir – The directory the run would write into.

• verbose – Whether to also log the plan.

Returns:

The plan – one step and one forward pass per prompt.

class entoscope.circuit.collector.CircuitCollector(operation_input: CircuitInput, model: TransformerBridge | CachedModel, save_dir: str | None = None)

Bases: AbstractCollector[CircuitInput]

Writes every prompt’s graphs, as GraphML and node-link JSON, and a prompts.json summary.

This base collector writes graph data only, and CircuitVisualizationCollector also draws a page per prompt.

Parameters:

• operation_input – The input the run iterates.

• model – The model the run reads.

• save_dir – The directory save writes into, or None to write nothing.

property result: list[PromptCircuit]

One circuit per prompt collected so far, in prompt order.

collect(circuit: PromptCircuit, index: int) → None

Take one prompt’s circuit.

Parameters:

• circuit – The circuit.

• index – The prompt’s position among the input’s prompts, collected in order.

output_files() → list[str]

List the file names save writes into save_dir, which a plan reports as the run’s output.

Subclasses extend this to mirror their own _save_result. It is deliberately not abstract: collector_class is an extension point, so a collector that declares no files keeps working and is reported as writing none. A plan asks a collector nothing was collected into, so an override may read the input but never the collected result.

Returns:

The file names, relative to save_dir.

class entoscope.circuit.collector.CircuitVisualizationCollector(operation_input: CircuitInput, model: TransformerBridge | CachedModel, save_dir: str | None = None)

Bases: CircuitCollector, FigureMixin[CircuitFigure]

A CircuitCollector that also draws each prompt’s circuit as circuit_{i}.html.

Kept separate from the base collector so that anything which only needs graph data draws no figure, and so visualization is opt-in by collector choice rather than a config flag.

property figures: list[dict[str, Figure]]

The plots each prompt’s circuit_{i}.html holds, drawn from what the run collected so far.

One dict per prompt, in prompt order, keyed routes then graph. They are drawn afresh on every read.

output_files() → list[str]

List the file names save writes into save_dir, which a plan reports as the run’s output.

Subclasses extend this to mirror their own _save_result. It is deliberately not abstract: collector_class is an extension point, so a collector that declares no files keeps working and is reported as writing none. A plan asks a collector nothing was collected into, so an override may read the input but never the collected result.

Returns:

The file names, relative to save_dir.

figure_class

alias of CircuitFigure

Interactive Plotly rendering of a discovered circuit.

Lays the graph out as the Information Flow Routes figures do: token position on the x axis, model depth on the y axis (input tokens at the bottom, the final layer at the top). Attention edges are drawn per head: the coloured midpoint markers encode which head carries each edge (colour) and how strong it is (size). One self-contained HTML file is produced per prompt, holding two figures: the pruned routes-to-prediction (headline) and the full contribution graph.

class entoscope.circuit.visualization.CircuitFigure(circuit: PromptCircuit, model: TransformerBridge | CachedModel, config: CircuitConfig)

Bases: Figure[str]

The page of one prompt’s circuit: the routes to the prediction, then the full contribution graph.

Each panel is named by the PromptCircuit attribute it draws, routes then graph.

Parameters:

• circuit – The prompt’s circuit.

• model – The model the circuit was discovered in.

• config – The run’s config.

property title: str

The page’s title.

panels() → list[str]

Name the page’s panels.

Returns:

The panels, in page order, each drawn as its own plot.

figure(panel: str) → Figure

Draw one panel.

Parameters:

panel – One of panels().

Returns:

The panel’s plot.

class entoscope.circuit.method.circuit.PromptCircuit(prompt: str, str_tokens: list[str], graph: MultiDiGraph, routes: MultiDiGraph)

Bases: object

One prompt’s discovered circuit.

The circuit page and the GraphML files read the attributes Information Flow Routes writes: each node a kind (input, after_attn, mlp or after_ffn), layer, position and token, and each edge a kind (attention, residual or mlp), weight and head (-1 off attention). A method writing other attributes runs with CircuitCollector, which draws no page.

Variables:

• prompt (str) – The prompt.

• str_tokens (list[str]) – The prompt’s token strings, one per position.

• graph (networkx.classes.multidigraph.MultiDiGraph) – The full thresholded contribution graph, with a separate attention edge per head.

• routes (networkx.classes.multidigraph.MultiDiGraph) – The subgraph of the edges on paths leading to the traced prediction positions.

class entoscope.circuit.method.circuit.AbstractCircuitMethod(model: TransformerBridge | CachedModel, config: CircuitConfig)

Bases: ABC, Generic[CircuitMethodConfigT]

The base of every circuit discovery method: discovers the circuit behind one prompt at a time.

A method is built once per run. Its type argument names the dataclass of its method config, extending CircuitMethodConfig, which a config file’s method is read into. A method naming none takes CircuitMethodConfig itself, which no field fits.

Parameters:

• model – The model the method reads.

• config – The run’s config.

Variables:

method_config – The run’s method config, typed by the method’s type argument.

abstractmethod build(prompt: str, tokens: Int[Tensor, 'batch pos']) → PromptCircuit

Discover the circuit behind one prompt.

Parameters:

• prompt – The prompt.

• tokens – The prompt’s tokens, [1, pos].

Returns:

The prompt’s circuit.

activation_names() → list[str] | None

Name the transformer activations build caches, which a plan reports.

Returns:

The transformer activation names, or None when build caches every activation, the default.

Information Flow Routes circuit-discovery method.

Builds a per-head contribution graph over a single forward pass and prunes the routes that lead to the model’s prediction, following Ferrando & Voita, “Information Flow Routes: Automatically Interpreting Language Models at Scale” (2024, https://arxiv.org/abs/2403.00824).

The node scheme (X0_/A/M/I) and the per-layer contribution loop follow the LM Transparency Tool’s llm_transparency_tool/routes/graph.py, but the graph here is a networkx.MultiDiGraph so that each attention head is its own edge (the upstream GraphBuilder uses a DiGraph and sums attention over heads). Route pruning is reimplemented to work on a multigraph (the upstream helper is DiGraph-only: its filter_edge lambda takes two args, which breaks for multigraphs).

The numerical contributions come from the vendored entoscope.circuit.vendor.contributions module, copied from the LM Transparency Tool under CC BY-NC 4.0 (see entoscope/circuit/vendor/NOTICE). Because that code is non-commercial, this method, and any results it produces, may be used for non-commercial purposes only.

class entoscope.circuit.method.information_flow_routes.NodeType(*values)

Bases: Enum

A residual-stream node component.

Node names are {prefix}{layer}_{token}. INPUT is always built at layer 0, so it renders as X0_{token}.

Variables:

• prefix – The node-name prefix.

• kind – The node’s kind attribute value.

node_name(layer: int, token: int) → str

Name the node of this component at a layer and position.

Parameters:

• layer – The node’s layer.

• token – The node’s position.

Returns:

The node’s name, such as A3_5.

entoscope.circuit.method.information_flow_routes.decomposed_attn(model: TransformerBridge | CachedModel, cache: ActivationCache, layer: int, batch_i: int = 0) → Float[Tensor, 'pos key_pos head d_model']

Split one layer’s attention output by query position, key position and head.

Reconstructs each head’s output written to the residual stream, split by the key position it came from, before the residual add. No layernorm is applied, since the attention output is added to the residual stream unnormalized. Summing over key_pos and head recovers attn_out minus the b_O bias. Grouped-query attention is handled by repeating the value heads across their group.

Parameters:

• model – The model.

• cache – The forward pass’s activations, holding the layer’s values and pattern.

• layer – The layer.

• batch_i – The batch row to read.

Returns:

Each head’s contribution from each key position to each query position, in d_model space.

entoscope.circuit.method.information_flow_routes.build_information_flow_graph(model: TransformerBridge | CachedModel, cache: ActivationCache, n_tokens: int, str_tokens: list[str], threshold: float) → MultiDiGraph

Build the full per-head contribution graph of one cached forward pass, at batch row 0.

Parameters:

• model – The model.

• cache – The forward pass’s activations, holding every activation READ_ACTIVATIONS names.

• n_tokens – The prompt’s token count.

• str_tokens – The prompt’s token strings, which label the nodes.

• threshold – The renormalizing edge threshold: contributions below it are dropped and the rest renormalized.

Returns:

The graph, with one attention edge per head.

Raises:

NotImplementedError – If the model runs attention and MLP in parallel.

entoscope.circuit.method.information_flow_routes.resolve_positions(positions: list[int], n_tokens: int) → list[int]

Resolve prediction positions, negative ones counting from the end of the prompt.

Parameters:

• positions – The positions.

• n_tokens – The prompt’s token count.

Returns:

The distinct positions, counted from the start, in the order given.

Raises:

ValueError – If a position is out of range for the prompt.

entoscope.circuit.method.information_flow_routes.prune_routes_to_predictions(graph: MultiDiGraph, n_layers: int, starting_positions: list[int], threshold: float) → MultiDiGraph

Keep only the edges on paths that reach the output nodes, traced backwards over strong edges.

Mirrors the upstream build_paths_to_predictions (reverse the graph, restrict to edges above threshold, DFS back from each output node) but operates on a MultiDiGraph and merges the per-position trees into a single graph.

Parameters:

• graph – The full contribution graph.

• n_layers – The model’s layer count, which locates the output nodes.

• starting_positions – The positions whose prediction the routes lead to.

• threshold – The weight an edge must exceed to be followed.

Returns:

The routes, empty when no strong edge reaches an output node.

class entoscope.circuit.method.information_flow_routes.InformationFlowRoutesConfig(*, type: str, threshold: float = 0.04, path_threshold: float | None = None, prediction_positions: list[int] = <factory>)

Bases: CircuitMethodConfig

The method config of information_flow_routes.

Variables:

• threshold (float) – The renormalizing edge threshold the full graph is built with. Contributions below it are dropped and the rest renormalized, so the incoming weights of each representation still sum to 1, as in the Information Flow Routes paper. 0.0 keeps every edge.

• path_threshold (float | None) – The threshold the routes to the prediction are pruned with, threshold itself when None.

• prediction_positions (list[int]) – The positions whose prediction the routes are traced back from, negative ones counting from the end of each prompt. The default [-1] is the next-token prediction.

resolved_path_threshold() → float

Resolve the threshold the routes are pruned with.

Returns:

path_threshold, or threshold when it is unset.

class entoscope.circuit.method.information_flow_routes.InformationFlowRoutes(model: TransformerBridge | CachedModel, config: CircuitConfig)

Bases: AbstractCircuitMethod[InformationFlowRoutesConfig]

Information Flow Routes (Ferrando and Voita, 2024): a per-head contribution graph and its routes.

For non-commercial use only, as it builds on CC BY-NC 4.0 code.

build(prompt: str, tokens: Int[Tensor, 'batch pos']) → PromptCircuit

Build one prompt’s contribution graph from a single forward pass, and prune its routes.

Parameters:

• prompt – The prompt.

• tokens – The prompt’s tokens, [1, pos].

Returns:

The prompt’s circuit.

Raises:

• NotImplementedError – If the model runs attention and MLP in parallel.

• ValueError – If a prediction position is out of range for the prompt.

activation_names() → list[str]

Name the transformer activations build caches.

They are the residual stream, the MLP output, and the attention values and pattern.

Returns:

The transformer activation names, layer by layer.

class entoscope.circuit.method.factory.CircuitMethodFactory

Bases: RegistryFactory[AbstractCircuitMethod]

Resolves the circuit discovery method a config’s method.type names and builds it for a run.

It holds information_flow_routes and every method registered with register.

classmethod create(model: TransformerBridge | CachedModel, config: CircuitConfig) → AbstractCircuitMethod

Build the method a config names, once per run. The method then builds one prompt’s circuit at a time.

Parameters:

• model – The model the method reads.

• config – The run’s config, whose method.type names the method.

Returns:

The method.

Raises:

ValueError – If no method is registered under the name, or the method rejects the config.