entoscope.operation

The shared machinery every operation is built from.

Its classes are importable from entoscope.operation, and ConfigFactory from entoscope too. from_dict and to_dict, which read a config from its raw shape and write it back, are in entoscope.util.data.

class entoscope.operation.config.AbstractConfig(*, config_id: str | None = None)

Bases: object

The base of every operation’s config.

Variables:

config_id (str | None) – The id of the YAML entry the config was read from, set by ConfigFactory. It titles the run’s plan and result pages, and a config built in code may leave it unset.

save(save_dir: str) → None

Write the config as config.json beside a run’s results.

Parameters:

save_dir – The directory the run’s results are written to.

Raises:

TypeError – If a field holds a value with no JSON form, such as a function or a tensor.

class entoscope.operation.config.ConfigFactory(config_path: str)

Bases: Generic[T]

Reads configs out of a YAML file of config_id: config entries.

The type argument is the class an entry is read into, so it is required: ConfigFactory[LensConfig]("configs.yaml").create("my_lens"). Reading is strict, so an unknown or missing key raises, and the file is read once per factory. A value may name environment variables as ${NAME}: each is replaced by the variable’s value, or by NAME itself when it is unset, and a value that is a single variable holding digits becomes an int.

Parameters:

config_path – The YAML file to read.

create(config_id: str) → T

Read one entry of the file into the factory’s config class.

Parameters:

config_id – The entry to read, which also becomes the config’s config_id.

Returns:

The config the entry describes.

Raises:

• KeyError – If the file has no entry config_id.

• ValueError – If the entry has a key the config does not, lacks a required one, or holds a value its field cannot take.

The method config a config runs its method with, and how a method names its class.

class entoscope.operation.method.MethodConfig(*, type: str)

Bases: object

The base of every method config: the method a config runs, and that method’s settings.

Each operation’s method config base extends this one and names its factory. type has no default because one method config class serves every method that takes it (each method with no fields of its own takes its operation’s base), so only the instance can say which method it configures.

Variables:

type (str) – The name the method is registered under.

classmethod factory() → type[RegistryFactory[Any]]

Name the factory the operation’s methods register with.

Returns:

The operation’s method factory.

Raises:

NotImplementedError – If the class names no factory.

classmethod from_config(value: Any) → MethodConfig

Read a config file’s method into the method config class its method takes.

Parameters:

value – The method’s name alone, or a mapping of type and the method’s fields.

Returns:

An instance of the config class the method named by type declares.

Raises:

ValueError – If value is neither a name nor a mapping, names no type, names a method nothing registered, or has a key the method config does not.

entoscope.operation.method.method_config_class(method: type) → type[MethodConfig]

Find the config class a method or a dataset names as its first type argument.

A TypeVar left open gives its bound, and a class naming none takes that of the class it extends. Type arguments are erased at runtime, so the subscripted base in a class statement is their only record.

Parameters:

method – A method or dataset class.

Returns:

The config class it is configured with.

Raises:

TypeError – If neither the class nor any class it extends names a config class.

Name-to-class registries: the extension points a user fills, a dataset or an operation’s method.

The built-ins register through the same register a user calls, so a name entoscope ships with differs from a user’s only in having been taken first.

class entoscope.operation.registry.RegistryFactory

Bases: Generic[T]

Resolves a name to the class registered under it, which extends the class the factory names.

A subclass names that class as its type argument and what an implementation is called in messages as kind, and each keeps its own entries. A name belongs to one implementation, so a typo or a second implementation cannot shadow the first one, while defining the same class again replaces it, which is what re-running a notebook cell does.

Variables:

kind (ClassVar[str]) – What an implementation is called in messages, such as "lens method".

classmethod register(name: str) → Callable[[ImplementationT], ImplementationT]

Register the decorated class under name.

The class is checked as it is registered rather than at its first run.

Parameters:

name – The name a config’s type selects the class by.

Returns:

A class decorator that registers the class and returns it unchanged.

Raises:

• ValueError – If name is empty or already taken by another class.

• TypeError – If the decorated class does not extend the factory’s base class, or leaves one of its abstract methods unimplemented.

classmethod get(name: str) → type[T]

Look up the class registered under name.

Parameters:

name – A registered name.

Returns:

The class registered under name.

Raises:

ValueError – If nothing is registered under name. The message lists the registered names.

classmethod registered() → list[str]

List the registered names.

Returns:

Every name registered with this factory, sorted.

class entoscope.operation.input.OperationInput

Bases: ABC

The base of every operation’s input: a config together with the prompts it runs on.

An input fills itself lazily through prepare, which resolves what a run iterates and then gathers the activations it reads. Each of the two is done once, so plan followed by apply on the same input resolves the index once.

Variables:

config (entoscope.operation.config.AbstractConfig) – The config the input runs.

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

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.

abstractmethod to(device: str | device = 'cpu') → Self

Move the tensors the input holds to device.

Parameters:

device – The device to move to.

Returns:

The input itself, moved in place.

class entoscope.operation.operator.Operator

Bases: ABC

Runs an operation’s method over its input and hands every result to a collector.

An operator is built from (model, collector_class=...). The collector class decides what a run leaves behind, which is how a run opts in or out of its HTML page. The method is the one the input’s config names, built when the run starts.

abstractmethod apply(operation_input: OperationInput, save_dir: str | None = None, return_collector: bool = False, verbose: bool = True) → AbstractCollector | None

Run the operation and save what it collected into save_dir.

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:

• operation_input – The config and prompts to run, prepared in place as the run needs.

• save_dir – The directory to write the run’s files 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 what the run collected, if return_collector is set, and None otherwise.

abstractmethod plan(operation_input: OperationInput, save_dir: str | None = None, verbose: bool = True) → OperationPlan

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

The input’s index is resolved exactly as the run resolves it, and kept: passing the same input to apply afterwards does not resolve it again.

Parameters:

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

• save_dir – The directory the run would write into, named in the plan’s outputs.

• verbose – Whether to also log the plan.

Returns:

The plan – the run’s steps, forward and backward passes, activations, outputs and index.

class entoscope.operation.plan.OperationPlan(operation: str, method: str, config: AbstractConfig, steps: int, steps_label: str, forward_passes: int, backward_passes: int = 0, batches: int | None = None, index: list[DataFrame] = <factory>, prompts: list[tuple[str, str, int]]=<factory>, activations: list[str] = <factory>, reads_full_cache: bool = False, save_dir: str | None = None, outputs: list[str] = <factory>)

Bases: object

What a run would do, as an operator’s plan reports it.

str(plan) gives describe().

Variables:

• operation (str) – The operation’s name, such as "Lens".

• method (str) – The name of the method the config runs.

• config (entoscope.operation.config.AbstractConfig) – The config the run would be made under.

• steps (int) – How many units of work the run makes.

• steps_label (str) – What a unit of work is: an index element, a probe, a patched run or a prompt.

• forward_passes (int) – Every forward pass the run makes, the cached runs of the prompts included.

• backward_passes (int) – Every backward pass the run makes.

• batches (int | None) – How many batches the run makes, set only for a method that runs the model a batch at a time.

• index (list[pandas.DataFrame]) – One frame per index, the same frame index.csv would hold.

• prompts (list[tuple[str, str, int]]) – One (label, prompt, token count) per prompt.

• activations (list[str]) – The transformer activation names the run caches.

• reads_full_cache (bool) – Whether the run caches every activation, in which case activations is empty.

• save_dir (str | None) – The directory the run would write into.

• outputs (list[str]) – The file names the run would write into save_dir.

describe(max_rows: int = 100) → str

Render the plan as text, which is what str(plan) and the log show.

Parameters:

max_rows – The most rows of each index to print.

Returns:

The plan, one fact per line, followed by each index.

class entoscope.operation.collector.AbstractCollector(operation_input: InputT, model: TransformerBridge | CachedModel, save_dir: str | None = None)

Bases: ABC, Generic[InputT]

Accumulates one operation run’s per-location results and writes them to save_dir.

A subclass owns the shape of a result (a DataFrame, a list of graphs, …), which it allocates as it is constructed, fills one location’s cell at a time through collect and exposes as result. save creates save_dir and hands it to _save_result, which owns everything the run leaves behind: its result, the index files beside it and the config it ran under.

Every collector takes the model whether or not it reads one, so an operator constructs whichever collector_class it was given the same way, and swapping in a collector that labels the prompt’s tokens (or reads cfg) for one that does not is only a class name.

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.

abstract property result: Any

What the run collected so far.

abstractmethod collect(*args: Any, **kwargs: Any) → None

Take one result from the run’s method, filed under the index position it belongs to.

Parameters:

• *args – The result and its position, as each collector declares them.

• **kwargs – The same, by keyword.

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.

save() → None

Write the run’s files into save_dir, created if missing, and log where they went.

Writes nothing when save_dir is None.

class entoscope.operation.collector.FrameCollector(operation_input: InputT, model: TransformerBridge | CachedModel, save_dir: str | None = None)

Bases: AbstractCollector[InputT]

A collector whose result is a frame with one row per row of the index the run swept.

Cells arrive one at a time under the index position they fill and are glued back onto that index’s own columns, so a subclass says only which indices the run has and what a cell holds. It writes result.csv, an index CSV per index and config.json.

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: DataFrame

The index’s columns followed by the collected cells, one row per index row.

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.operation.collector.FigureMixin

Bases: Generic[FigureT]

Names the figure a collector draws, read off the collector’s own type argument.

The figure is the last type parameter of whichever FigureMixin a collector names, so it is stated once, in the base list. An explicit assignment still wins, and an intermediate subclass that leaves the parameter free resolves nothing and inherits.

Variables:

figure_class (type[entoscope.operation.collector.FigureT]) – The Figure the collector draws.

class entoscope.operation.collector.VisualizationCollector(operation_input: InputT, model: TransformerBridge | CachedModel, save_dir: str | None = None)

Bases: FrameCollector[InputT], FigureMixin[FigureT], ABC

A FrameCollector that also draws its frame as an interactive Plotly result.html.

Kept apart from the base collectors so that anything which only needs the CSVs builds no figure, and so visualization is opt-in by collector choice rather than by a config flag. Every GridFigure is built from the same four arguments and reads what it draws off them, so a subclass names its figure in its base list and nothing else. The page’s plots are handed back as figures too, so a run that writes nothing can still be looked at.

property figures: dict[Any, go.Figure]

The plots result.html holds, drawn from what the run collected so far.

Keyed by the figure’s panels, in page order. 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.

Index-grid plotting shared by the operations whose result is a tidy frame over an Index.

An Index lays its rows out as dense per-axis codes and the result frame carries the original source_*/target_* values, so laying one out is the same problem for every operation and only the cell payload differs. Everything produced here ends up inside plotly HTML, so values are escaped once, in format_value, rather than at each call site. Figure, the page every operation’s HTML is written through, grid-drawn or not, lives here too.

A run rarely reads one number off itself: patching names several metrics and the lens keeps a whole top-k. A page is therefore a list of panels, each drawn as its own plot from the same result, and an operation says only what its panels are.

entoscope.operation.visualization.format_value(value: Any) → str

Format a result cell for a page, HTML-escaped.

Parameters:

value – An enum, a list or tuple, a float, None or anything str can print.

Returns:

An enum’s value, a list in brackets, a float to four significant digits, "-" for a missing value, or str(value).

entoscope.operation.visualization.label_column(frame: DataFrame, axis: str, prefixes: Sequence[str]) → str | None

Find the source_* or target_* column holding the real values behind axis’s dense codes.

A column that is not fully populated is rejected: concatenating a shorter source index onto a longer target one leaves NaN behind, and a half-filled axis labels nothing.

Parameters:

• frame – An index’s frame, as Index.to_df lays it out.

• axis – An axis of the index, such as "layer".

• prefixes – The column prefixes to try, in order of preference.

Returns:

The first fully populated column, or None when there is none.

entoscope.operation.visualization.tick_text(frame: DataFrame, grid: IndexGrid, axis: str, prefixes: Sequence[str], tokens: dict[str, list[str]]) → dict[int, str]

Label every code of one axis of the grid.

A position is labeled with its token, a row without a value on the axis with its full coordinates.

Parameters:

• frame – The index’s own frame, one row per result row, as Index.to_df lays it out.

• grid – The grid.

• axis – The axis to label.

• prefixes – The column prefixes to read values from, in order of preference.

• tokens – The prompt’s token strings, keyed by the position column they label.

Returns:

The label of each of the axis’s codes.

entoscope.operation.visualization.describe_row(row: Series, prefix: str) → str

Describe the location of one index row on one side of the index.

Parameters:

• row – A row of the index’s frame, as Index.to_df lays it out.

• prefix – The side to describe, "source_" or "target_".

Returns:

The activation followed by every axis the row names, such as attn.z · layer 3 · head 5.

entoscope.operation.visualization.diverging_range(values: Sequence[float], midpoint: float) → tuple[float, float]

Centre a colour range on midpoint wide enough for every value.

Parameters:

• values – The values to colour, NaN and infinities ignored, as no finite range holds them.

• midpoint – The value drawn in the middle colour.

Returns:

The range’s low and high ends, one unit either side of midpoint when every value equals it.

entoscope.operation.visualization.write_html(path: str, title: str, *figures: Figure) → None

Write figures one under the other into a self-contained HTML page.

plotly.js is embedded once, with the first figure, so the page opens with no network.

Parameters:

• path – The file to write.

• title – The page’s title.

• *figures – The figures, in page order.

class entoscope.operation.visualization.Figure

Bases: ABC, Generic[PanelT]

A self-contained interactive HTML page, one plotly figure per panel.

title names the page, and each panel titles its own plot, so the two do not have to repeat.

abstract property title: str

The page’s title.

abstractmethod panels() → list[PanelT]

Name the page’s panels.

Returns:

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

abstractmethod figure(panel: PanelT) → Figure

Draw one panel.

Parameters:

panel – One of panels().

Returns:

The panel’s plot.

figures() → dict[PanelT, Figure]

Draw every panel.

Returns:

Each panel’s plot, keyed by the panel, in page order.

render(path: str) → None

Write the page.

Parameters:

path – The HTML file to write.

class entoscope.operation.visualization.IndexGrid(x_axis: str | None, y_axes: list[str], codes: dict[str, list[int]], x_codes: list[int], y_keys: list[tuple[int, ...]], rows: dict[tuple[tuple[int, ...], int], int])

Bases: object

Where each result row sits once the axes a run varies over are laid out.

Built from the index’s own layout codes rather than from the result frame, so the axis rule can be read, and tested, without a run behind it. An axis the run does not vary over is dropped: it offers no comparison and belongs in the subtitle. Axes beyond the first two fold into composite y keys rather than being aggregated away, so nothing is hidden. y_keys ascends and plotly draws the first at the bottom, putting layer 0 there as the circuit figures do.

Variables:

• x_axis (str | None) – The axis across, the last varying one in AXIS_ORDER, or None when nothing varies.

• y_axes (list[str]) – The other varying axes, folded together into the rows.

• codes (dict[str, list[int]]) – Every axis’s code for every result row.

• x_codes (list[int]) – The distinct codes across, ascending.

• y_keys (list[tuple[int, ...]]) – The distinct row keys, one code per y_axes entry, ascending.

• rows (dict[tuple[tuple[int, ...], int], int]) – The result row at each (y_key, x_code) cell.

classmethod of(layout: DataFrame) → IndexGrid

Lay out an index.

Parameters:

layout – The index’s layout codes, one column per axis, as Index.layout_df holds them.

Returns:

The grid.

property is_bar: bool

Whether at most one axis varies, so the grid is drawn as a bar chart.

lay_out(values: Sequence[Any], fill: Any = None) → list[list[Any]]

Place one value per result row onto the grid.

Parameters:

• values – One value per result row.

• fill – The value of a cell no row lands on.

Returns:

One list per row key, each holding one value per x code.

entoscope.operation.visualization.x_labels(frame: DataFrame, grid: IndexGrid, prefixes: Sequence[str], tokens: dict[str, list[str]]) → list[str]

Label the grid’s columns.

Parameters:

• frame – The index’s own frame, one row per result row, as Index.to_df lays it out.

• grid – The grid.

• prefixes – The column prefixes to read values from, in order of preference.

• tokens – The prompt’s token strings, keyed by the position column they label.

Returns:

One label per x code, or the single row’s location when no axis varies.

entoscope.operation.visualization.y_labels(frame: DataFrame, grid: IndexGrid, prefixes: Sequence[str], tokens: dict[str, list[str]]) → list[str]

Label the grid’s rows.

Parameters:

• frame – The index’s own frame, one row per result row, as Index.to_df lays it out.

• grid – The grid.

• prefixes – The column prefixes to read values from, in order of preference.

• tokens – The prompt’s token strings, keyed by the position column they label.

Returns:

One label per row key, each axis named when several fold into the rows.

entoscope.operation.visualization.cell_labels(frame: DataFrame, grid: IndexGrid, prefixes: Sequence[str], tokens: dict[str, list[str]]) → list[str]

Give every result row the coordinates of its cell.

Parameters:

• frame – The index’s own frame, one row per result row, as Index.to_df lays it out.

• grid – The grid.

• prefixes – The column prefixes to read values from, in order of preference.

• tokens – The prompt’s token strings, keyed by the position column they label.

Returns:

One label per result row, such as layer 3 · position 4: " France".

entoscope.operation.visualization.titled(kind: str, config: AbstractConfig) → str

Title a page after the operation and, when it has one, the config’s id.

Parameters:

• kind – What the page shows, such as "Logit lens".

• config – The config the run was made under.

Returns:

kind, followed by the config’s id in parentheses when it has one.

class entoscope.operation.visualization.GridFigure(result: DataFrame, operation_input: InputT, model: TransformerBridge | CachedModel, index: Index)

Bases: Figure[PanelT], Generic[InputT, PanelT]

An operation’s result frame drawn as the grid its index describes, once per panel.

Choosing the axes, pivoting the rows onto them and falling back to a bar when only one of them varies is the same work whatever was run, so it lives here. A subclass says only what its own result means: which column a panel holds its value in, how to colour it, and what a cell reads on hover. Those answers are read off the run’s input, which comes along for that reason, and because it lets a collector name a figure class and construct it without knowing what produced the frame. Every panel shares one grid and one set of axis labels, so those are resolved once, at construction.

Parameters:

• result – The run’s result frame.

• operation_input – The input the run iterated.

• model – The model the run read.

• index – The index that laid the result out, whose layout codes build the grid.

Variables:

• coordinates – The index’s own frame, which every cell, tick and hover is labelled from: the result frame also holds the values, and a value column named with a side’s prefix would read as a coordinate.

• prefixes – The column prefixes axis values are read from, in order of preference: the index’s own layout_prefixes, so the axes are labelled by the side the grid was laid out along.

• colorscale – The plotly colour scale cells are drawn with.

abstractmethod heading(panel: PanelT) → str

Name one panel’s plot.

Parameters:

panel – One of panels().

Returns:

The plot’s heading, shown above the subtitle line.

abstractmethod subtitle(panel: PanelT) → str

Describe what one panel’s plot holds, under its heading.

Parameters:

panel – One of panels().

Returns:

The subtitle, as HTML.

abstractmethod value_column(panel: PanelT) → str

Name the result column one panel colours its cells by.

Parameters:

panel – One of panels().

Returns:

A column of the result frame.

abstractmethod hover(panel: PanelT, position: int, coordinate: str) → str

Write what one cell shows on hover.

Parameters:

• panel – One of panels().

• position – The result row the cell is built from.

• coordinate – The cell’s laid-out coordinates, from cell_labels.

Returns:

The hover text, as HTML.

value_title(panel: PanelT) → str

Title one panel’s colour bar.

Parameters:

panel – One of panels().

Returns:

The title, as text, by default the value column’s name.

value_range(panel: PanelT, values: Sequence[float]) → tuple[float, float]

Choose the range one panel’s colours span.

Parameters:

• panel – One of panels().

• values – The panel’s values, one per result row.

Returns:

The low and high ends of the colour range, by default 0 and 1.

bar_yaxis(panel: PanelT) → dict

Lay out the value axis of one panel drawn as a bar chart.

Parameters:

panel – One of panels().

Returns:

plotly yaxis settings, by default only the axis title.

cell_text(panel: PanelT) → list[str] | None

Choose the label printed inside each cell of one panel.

The size cap is applied by figure, so a subclass answers for its own result alone.

Parameters:

panel – One of panels().

Returns:

One label per result row, or None to leave the cells bare, the default.

str_tokens() → dict[str, list[str]]

Read the prompt’s token strings, to label the position axes with.

Read once at construction. A figure whose index carries no position, or a run with no prompt to read one from, labels its ticks with the coordinates alone.

Returns:

The token strings for each position column of the result, keyed by that column. Empty by default.

figures() → dict[PanelT, Figure]

Draw every panel, or only the first, saying so, when the index has no rows.

Returns:

Each panel’s plot, keyed by the panel, in page order.

figure(panel: PanelT) → Figure

Draw one panel as a heatmap over the grid, or a bar chart when at most one axis varies.

Parameters:

panel – One of panels().

Returns:

The panel’s plot.