entoscope.model¶

Locating activations, the index operations iterate, and model access.

Model, CachedModel, ModelLocatorConfig, the locator transforms and the ModelAxis and ModelLayerLocation enums are importable from entoscope.model, and all but Model from entoscope too. The index classes are imported from the module defining them, entoscope.model.index.

entoscope.model.base.unwrap(model: TransformerBridge | CachedModel) → Any[source]¶

Remove a CachedModel wrapper.

Parameters:: model – A model, wrapped or not.
Returns:: The transformer itself.

entoscope.model.base.is_bridge(model: TransformerBridge | CachedModel) → bool[source]¶

Tell a TransformerBridge from a HookedTransformer.

Parameters:: model – A model, wrapped in a CachedModel or not.
Returns:: Whether the transformer is a TransformerBridge.

class entoscope.model.config.ModelAxis(*values)[source]¶

Bases: StrEnum

The axes an index row names its location on, and the row’s activation.

layer, head, neuron and position are axes a locator can enumerate, and so are dest_pos and src_pos, the query and the key position of attn.pattern (the names transformer_lens’s patching gives them). activation is not an axis but a key of every row: it is set once on the locator (or on each coordinate) and stamped on every row.

class entoscope.model.config.ModelLayerLocation(*values)[source]¶

Bases: StrEnum

The activations a locator can name, valued as a config file spells them.

Each maps onto a transformer activation name through TRANSFORMER_ACT_NAME, and the comment beside a member is its tensor’s axes. embed is not inside a block, so a locator reading it names no layer. attn.in and mlp.in are what the sublayer itself reads, after the norm before it on a model that has one, and attn.out and mlp.out what it adds to the residual stream, after the norm following it on a model that norms a sublayer’s output (Gemma 2 and 3 and OLMo 2, for instance). attn.out_pre_norm and mlp.out_pre_norm are the sublayer’s output before that norm, the same tensor as attn.out and mlp.out on any other model. Those four are hooks of a TransformerBridge’s sublayers, which a HookedTransformer does not have.

class entoscope.model.config.ModelLocationIndex(df: DataFrame = <factory>)[source]¶

Bases: object

A locator resolved against a model and a prompt length.

Variables:: df (pandas.DataFrame) – One row per location, one column per axis the locator names and one for the activation.

entoscope.model.config.count_from_end(value: int, max_value: int | None, axis: ModelAxis) → int[source]¶

Resolve a value on an axis, a negative one counting from the axis’s end.

Parameters:

value – A value on the axis.
max_value – The axis’s size. None is a prompt length that varies, so the value stays symbolic and torch resolves a negative one against each prompt’s own cache at slice time.
axis – The axis, which a message names.

Returns:

The value, counted from the start when the axis’s size is known.

Raises:

ValueError – If the axis’s size is known and the value lies past either end of it, which would name a row the model does not have.

Enumerate every combination of the axes’ values.

Parameters:

column_names – The axes.
column_ranges – Each axis’s values, or its size for every value from 0.

Returns:

One row per combination, one column per axis.

class entoscope.model.config.IndexTransform(*args, **kwargs)[source]¶

Bases: Protocol

A config section that rewrites the index DataFrame in the parse pipeline.

apply(index_df: DataFrame, location_to_max_value: dict) → DataFrame[source]¶

Rewrite the index.

Parameters:

index_df – The index so far.
location_to_max_value – The size of every axis the transform names, so a negative value counts from its end.

Returns:

The rewritten index.

class entoscope.model.config.Filter(inclusion_terms: list[dict[~entoscope.model.config.ModelAxis, int | dict | list]]=<factory>, exclusion_terms: list[dict[~entoscope.model.config.ModelAxis, int | dict | list]]=<factory>)[source]¶

Bases: object

Keeps or drops index rows by a list of patterns.

Each pattern maps an axis to a constraint: the keys of one pattern are AND-ed, the patterns are OR-ed, and {"not": <pattern>} drops the rows it matches instead. A constraint is one of:

an int, matching that value;
{"range": [start, end]}, matching start to end with both ends included;
a list of the above, matching any of them.

Negative numbers count from the end. The result is the union of the positive patterns (everything when there are none) minus the union of the not patterns.

Variables:

inclusion_terms (list[dict[entoscope.model.config.ModelAxis, int | dict | list]]) – The patterns a row is kept by, keyed by ModelAxis members.
exclusion_terms (list[dict[entoscope.model.config.ModelAxis, int | dict | list]]) – The patterns a row is dropped by, keyed by ModelAxis members.

classmethod from_config(terms: list) → Filter[source]¶

Read a config file’s filter.

Parameters:: terms – The patterns, {"not": <pattern>} for one that drops rows. A pattern’s keys are ModelAxis members or their values.
Returns:: The filter.
Raises:: ValueError – If a term is not a non-empty mapping, or a key names no axis.

to_config() → list[source]¶

Write the filter as a config file’s list of terms, which from_config reads back.

Returns:: The inclusion patterns, then each exclusion pattern as {"not": <pattern>}.

included(location: ModelAxis, max_value: int | None) → set[int][source]¶

Find the values of one axis some inclusion pattern can keep.

Parameters:

location – The axis.
max_value – The axis’s size, None for a prompt length that varies.

Returns:

Every value of the axis when a pattern leaves it unconstrained, otherwise the values the patterns name.

Raises:

ValueError – If every value would be kept on an axis whose size varies with the prompt.

excluded(location: ModelAxis, max_value: int | None) → set[int][source]¶

Find the values of one axis an exclusion pattern on that axis alone drops.

Parameters:

location – The axis.
max_value – The axis’s size, None for a prompt length that varies.

Returns:

The values dropped whatever the row’s other axes hold.

location_range(location: ModelAxis, max_value: int | None) → list[int][source]¶

Find the values of one axis worth enumerating before the filter is applied row by row.

Parameters:

location – The axis.
max_value – The axis’s size, None for a prompt length that varies.

Returns:

The included values minus the excluded ones, ascending.

apply(index_df: DataFrame, location_to_max_value: dict) → DataFrame[source]¶

Keep the rows matching an inclusion pattern and drop those matching an exclusion pattern.

Parameters:

index_df – The index.
location_to_max_value – The size of every axis, so a negative value counts from its end.

Returns:

The filtered index.

Raises:

ValueError – If a pattern names an axis the index does not have, or holds an invalid constraint.

class entoscope.model.config.Group[source]¶

Bases: dict[ModelAxis, list[int] | str]

Merges the listed values of an axis into one row that addresses them together.

A mapping from axis to the values to merge: {head: [6, 9]} turns the rows of heads 6 and 9 into a single row reading both heads, and "all" merges every value the axis has. The duplicate rows a merge leaves are dropped. Merge heads, neurons or positions: every layer is a transformer activation of its own, so a row naming several layers is refused when an operation resolves it.

apply(index_df: DataFrame, location_to_max_value: dict) → DataFrame[source]¶

Merge the listed values of every axis.

Parameters:

index_df – The index.
location_to_max_value – The size of every axis, so a negative value counts from its end.

Returns:

The index with each merged group in one row.

class entoscope.model.config.Incremental[source]¶

Bases: dict[ModelAxis, list[int]]

Follows every row with rows that add growing prefixes of the listed values to an axis.

A mapping from axis to the values to add: {head: [6, 9]} turns the row of head 0 into rows for heads [0], [0, 6] and [0, 6, 9]. An axis the index does not have yet is added, and then only the grown rows exist.

apply(index_df: DataFrame, location_to_max_value: dict) → DataFrame[source]¶

Grow every row along every listed axis.

Parameters:

index_df – The index.
location_to_max_value – The size of every axis, so a negative value counts from its end.

Returns:

The index with the grown rows after each row they grow from.

class entoscope.model.config.Combination(location: list[ModelAxis] = <factory>, max_elements: int = 9223372036854775807)[source]¶

Bases: object

Expands each row into one row per combination of the values it holds on each listed axis.

Combinations run from one value up to max_elements, so it follows a group: after group: {head: all} over heads 6, 9 and 10, max_elements: 2 gives the rows [6], [9], [10], [6, 9], [6, 10] and [9, 10].

Variables:

location (list[entoscope.model.config.ModelAxis]) – The axes to expand.
max_elements (int) – The most values a combination holds.

apply(index_df: DataFrame, location_to_max_value: dict) → DataFrame[source]¶

Expand every row along every listed axis.

Parameters:

index_df – The index, holding a list of values on every listed axis.
location_to_max_value – Unused, as the values are resolved already.

Returns:

The index with each row replaced by its combinations.

class entoscope.model.config.Coordinates(iterable=(), /)[source]¶

Bases: list[dict[ModelAxis, int | list[int] | ModelLayerLocation]]

Explicit coordinates, each a mapping from ModelAxis to a value.

A value is an int or a list of ints read together, and a coordinate’s activation a ModelLayerLocation.

classmethod from_config(coordinates: list) → Coordinates[source]¶

Read a config file’s coordinates.

Parameters:: coordinates – The coordinates, keyed by ModelAxis members or their values, with an activation a ModelLayerLocation member or its value.
Returns:: The coordinates, keyed by ModelAxis members.
Raises:: ValueError – If coordinates is not a list, a coordinate is not a non-empty mapping, or a key or an activation is not valid.

to_df(model: TransformerBridge | CachedModel, n_positions: int | None, activation: ModelLayerLocation | None = None) → DataFrame[source]¶

Resolve the coordinates into an index, negative values counting from the end of their axis.

Parameters:

model – The model the axes are sized by.
n_positions – The prompt’s length, None when it varies.
activation – The activation of a coordinate that names none.

Returns:

One row per coordinate, in order, one column per key.

Raises:

ValueError – If a value is neither an int nor a list of ints.

class entoscope.model.config.ModelLocatorConfig(*, activation: ModelLayerLocation | None = None, location: list[ModelAxis] = <factory>, filter: Filter = <factory>, group: Group = <factory>, incremental: Incremental = <factory>, combination: Combination = <factory>, coordinates: Coordinates | None = None)[source]¶

Bases: object

Locates model activations, either declaratively or explicitly.

Declaratively, location names the index axes and the full range of every axis is narrowed by filter and reshaped by the transforms. Explicitly, coordinates enumerates the exact points, every coordinate naming the same axes. The two modes are mutually exclusive. Negative values count from the end of their axis.

Variables:

activation (entoscope.model.config.ModelLayerLocation | None) – The activation every row reads. Left unset when every coordinate names its own.
location (list[entoscope.model.config.ModelAxis]) – The axes the index enumerates.
filter (entoscope.model.config.Filter) – Which rows of the enumerated grid to keep.
group (entoscope.model.config.Group) – The values of an axis to merge into one row.
incremental (entoscope.model.config.Incremental) – The values of an axis to add to every row one at a time.
combination (entoscope.model.config.Combination) – The axes to expand into every combination of their values.
coordinates (entoscope.model.config.Coordinates | None) – The exact rows, instead of location and the transforms.

Raises:

ValueError – If location names activation, coordinates is combined with location, filter or a transform, the coordinates name different axes, activation is set both on the locator and in the coordinates, a location is named with no activation to read it in, an axis is named that the activation’s tensor does not have, or group or combination reshapes an axis the index does not enumerate.

activations() → set[ModelLayerLocation][source]¶

Name the activations the locator’s rows read.

Returns:: The locator’s activation, or the one each coordinate names, empty when none is named.

parse(model: TransformerBridge | CachedModel, n_positions: int | None) → ModelLocationIndex[source]¶

Resolve the locator into an index for a prompt of n_positions tokens.

When the prompt length varies, as when probing prompts of unequal length, the values of every position axis stay symbolic: a non-negative one counts from the start and a negative one from each prompt’s own end, which torch resolves at slice time. Other axes never depend on the prompt length.

Parameters:

model – The model the axes are sized by.
n_positions – The prompt’s length, or None when it varies.

Returns:

The index, one row per location.

Raises:

ValueError – If n_positions is None and the locator would enumerate every value of a position axis, or a filter or a coordinate is not valid.

The transformer activation name each activation is cached under, and how a hook writing one hands it on.

entoscope.model.activation.transformer_activation(activation: ModelLayerLocation, layer: int | None) → str[source]¶

Name an activation as transformer_lens caches it.

Parameters:

activation – The activation.
layer – The block it is read in, None for embed.

Returns:

The transformer activation name, such as blocks.3.hook_resid_post.

entoscope.model.activation.check_transformer_activations(model: TransformerBridge | CachedModel, names: Iterable[str]) → None[source]¶

Check that the model computes every transformer activation a run reads or writes.

Which it computes depends on its backend and architecture: a HookedTransformer has none of a TransformerBridge’s sublayer inputs and outputs, a model without a gated MLP no mlp.pre_linear and a parallel block no resid.mid. A row of a block’s activation that names several layers, or none, names one without its block, which no model computes. An input checks its activation_names() as it resolves its index, so such a run is refused before any forward pass, and by its plan too, rather than failing in transformer_lens once it looks the activation up.

Parameters:

model – The model the run reads.
names – The transformer activation names the run reads or writes.

Raises:

ValueError – If the model computes no transformer activation of one of the names.

entoscope.model.activation.writable(name: str, tensor: Tensor) → Tensor[source]¶

Pick the tensor a hook writes into and hands on.

Any other activation is written in place, so the write reaches whatever else holds the tensor, which resid.mid needs: a bridge hooks it on the input of the norm before the MLP, and only an in-place write reaches the residual stream the block keeps beside that norm. The name is the one the hook is registered under, as transformer_activation gives it, not the hook point’s own: OLMo 2’s hook_resid_mid is mlp.hook_in by another name.

Parameters:

name – The transformer activation name the hook writes.
tensor – The tensor the hook is handed.

Returns:

A copy of tensor for a sublayer’s input, which only the sublayer reads, and tensor itself otherwise.

class entoscope.model.index.IndexElement(config: T, index_names: list[str], index: list[int | list[int] | str | ModelLayerLocation])[source]¶

Bases: Generic[T]

One side of an index row: its values by axis, read off a locator config.

slices turns the values into the slice of the activation tensor they address, and transformer_activation is the transformer activation name that tensor is cached under. A list value addresses several entries of its axis at once, lists on several axes every combination of their entries, and an axis the row does not name is read whole.

Parameters:

config – The locator the row was resolved from.
index_names – The row’s keys, ModelAxis values.
index – The row’s values, one per key.

slices(batch: int | slice = slice(None)) → tuple[slice | int | list[int], ...][source]¶

Slice the element’s activation.

The batch row is an argument rather than state on the element: zip_broadcast hands the same element to every row a one-row index is paired with, so a row stored on it would be shared. The slices are read and written with read_slices and write_slices: indexing a tensor with them directly would pair the lists of several axes.

Parameters:: batch – The batch rows to read.
Returns:: One index per axis of the activation tensor, in the order the tensor holds its axes.

read_slices(tensor: Tensor, batch: int | slice = slice(None)) → Tensor[source]¶

Read the element’s block of an activation of its kind.

Parameters:

tensor – The activation.
batch – The batch rows to read.

Returns:

The block of every combination of the element’s values, every axis an int does not drop kept where the tensor holds it.

read_unit_slices(weight: Tensor) → Tensor[source]¶

Read the element’s units off a weight that holds one d_model vector per unit of its activation.

The weight’s leading axes are the activation’s other than batch and its positions, as model_weight lays it out.

Parameters:: weight – The weight.
Returns:: The block of the units the element names, d_model last.

write_slices(tensor: Tensor, value: Tensor, batch: int | slice = slice(None)) → Tensor[source]¶

Write into the block read_slices reads, in the tensor writable picks.

The value is cast to the tensor’s dtype and device first, as assigning into a slice would: a write through index tensors takes only an exact match, and a constant such as value_replace’s is float32 whatever the model’s dtype.

Parameters:

tensor – The activation to write into.
value – The block, or anything that broadcasts to it.
batch – The batch rows to write.

Returns:

The tensor written into – tensor itself, or a copy of a sublayer’s input.

remaining_axes() → list[ModelAxis | str][source]¶

Name the activation axes that survive slices(), batch excluded.

An int index drops its axis while a slice or list keeps it.

Returns:: The surviving axes, in tensor order.

activation() → ModelLayerLocation[source]¶

Read the row’s activation, as a member whichever way the frame stores it (pandas 3 keeps a string).

Returns:: The activation, or None when the locator names none.

transformer_activation() → str[source]¶

Name the row’s transformer activation, the name its tensor is cached under.

Returns:: The transformer activation name, such as blocks.3.hook_resid_post.

layer() → int | None[source]¶

Read the row’s layer.

Returns:: The layer, or None when the row names no single layer.

get(location: ModelAxis | str) → int | list[int] | str | None[source]¶

Read the row’s value on one axis.

Parameters:: location – The axis, or activation.
Returns:: The value, or None when the row does not name the axis.

class entoscope.model.index.SingleIndex(output_position: int, source: IndexElement[T], target: IndexElement | None = None)[source]¶

Bases: Generic[T]

One row of an Index: where to read, where to write, and which result row it fills.

Positions are what Index iteration already hands out in order, so a run needs no coordinate of its own to find its row.

Parameters:

output_position – The row’s position in the index.
source – Where the row reads.
target – Where the row writes, for patching.

Variables:

output – The row’s position in the index, the key a collector files its cell under and the row it lands on in the result frame.

property source: IndexElement[T]¶: Where the row reads.

property target: IndexElement[T]¶: Where the row writes. Only a patching index has one.

class entoscope.model.index.Index(source_config: T, source: ModelLocationIndex, target_config: T | None = None, target: ModelLocationIndex | None = None)[source]¶

Bases: Generic[T]

The rows an operation iterates, one SingleIndex each, built once as the index is.

Patching adds a target to write into, paired with the source row by row, so the two have as many rows, or one of them a single row, which pairs with every row of the other.

Parameters:

source_config – The locator the source was resolved from.
source – The resolved source.
target_config – The locator the target was resolved from, for patching.
target – The resolved target, for patching.

Variables:

layout_df – One integer code per row on every axis a result page lays the rows out on.
layout_prefixes – The result columns’ prefixes, the side layout_df was built from first, so a page labels its axes with the side that spans them.

Raises:

ValueError – If the source and the target cannot be paired row by row.

activation_names() → list[str][source]¶

Name the transformer activations the rows read or write.

A source without an activation, a value a patch writes, reads none.

Returns:: The transformer activation names, sorted.

target_activation_names() → list[str][source]¶

Name the transformer activations the target rows write.

Returns:: The transformer activation names, sorted, none without a target.

to_df() → DataFrame[source]¶

Lay the index out as index.csv holds it.

Returns:: One row per index row, the source’s columns prefixed source_ and the target’s target_, a single-row side repeated along the other.

to_csv(path: str)[source]¶

Write to_df() as a CSV file.

Parameters:: path – The file to write.

class entoscope.model.slicer.ActivationSlicer[source]¶

Bases: ABC

Orders an index row’s values into a slice of one kind of activation tensor.

Each kind of activation holds its axes in its own order, and a slicer owns that order.

Variables:: AXES (tuple[entoscope.model.config.ModelAxis | str, ...]) – The tensor’s axes, in order.

abstractmethod static extract(index: defaultdict[str, slice | int | list[int]]) → tuple[slice | int | list[int], ...][source]¶

Order a row’s values into a slice of the tensor.

Parameters:: index – The row’s value on each axis, keyed by axis, a missing axis read whole.
Returns:: One index per axis, in AXES order.

static get(activation: ModelLayerLocation) → type[ActivationSlicer][source]¶

Pick the slicer of an activation.

Parameters:: activation – The activation.
Returns:: The slicer for the activation’s tensor layout.
Raises:: ValueError – If the activation is unknown.

class entoscope.model.slicer.LinearActivationSlicer[source]¶

Bases: ActivationSlicer

Slices a [batch, pos, neuron] activation, such as the residual stream, a sublayer’s input or an MLP’s.

static extract(index: defaultdict[str, slice | int | list[int]]) → tuple[slice | int | list[int], ...][source]¶

Order a row’s values into a slice of the tensor.

Parameters:: index – The row’s value on each axis, keyed by axis, a missing axis read whole.
Returns:: One index per axis, in AXES order.

class entoscope.model.slicer.HeadActivationSlicer[source]¶

Bases: ActivationSlicer

Slices a [batch, pos, head, neuron] activation: an attention head’s query, key, value or output.

static extract(index: defaultdict[str, slice | int | list[int]]) → tuple[slice | int | list[int], ...][source]¶

Order a row’s values into a slice of the tensor.

Parameters:: index – The row’s value on each axis, keyed by axis, a missing axis read whole.
Returns:: One index per axis, in AXES order.

class entoscope.model.slicer.PatternActivationSlicer[source]¶

Bases: ActivationSlicer

Slices a [batch, head, dest_pos, src_pos] attention pattern: the weight each query gives each key.

static extract(index: defaultdict[str, slice | int | list[int]]) → tuple[slice | int | list[int], ...][source]¶

Order a row’s values into a slice of the tensor.

Parameters:: index – The row’s value on each axis, keyed by axis, a missing axis read whole.
Returns:: One index per axis, in AXES order.

entoscope.model.slicer.located_axes(activation: ModelLayerLocation) → set[ModelAxis][source]¶

Name the axes a locator may name for an activation.

Those are the axes of the activation’s tensor, batch aside, and the layer for an activation inside a block, which every activation but embed is.

Parameters:: activation – The activation.
Returns:: The axes.

Local-storage caching of model activations.

CachedModel is a transparent wrapper around a transformer_lens model, a TransformerBridge or a HookedTransformer, that memorizes run_with_cache to disk. Operations receive it in place of the raw model and never know the difference: every attribute other than run_with_cache is forwarded to the wrapped model, so tokenization, run_with_hooks (patching), weight access and cfg behave exactly as before, while repeated forward passes for the same (model, tokens, names_filter) are read back from cache_dir instead of recomputed. A forwarded method that hands back the model itself, as to and eval do, hands back the wrapper, so model = model.to(device) keeps the cache.

Usage:

model = TransformerBridge.boot_transformers(config.model_name)
model = CachedModel(model, cache_dir="./.activation_cache")  # cache_dir=None => pass-through
Lens(model).apply(LensInput(LensConfig(...), "Paris is located in"))

class entoscope.model.cache.ActivationDiskCache(cache_dir: str)[source]¶

Bases: object

Content-addressed on-disk store for (logits, activation cache) pairs.

Pure filesystem and key logic with no dependency on a live model, so the key derivation can be tested directly. Entries live at <cache_dir>/<model_name>/<digest>.pt and are written atomically, through a temporary file and os.replace, so an interrupted run never leaves a half-written file.

Parameters:: cache_dir – The directory entries are stored under.

key(*, model_fingerprint: str, model_args: tuple[Any, ...], names_filter: str | list[str], extra: dict[str, Any]) → str[source]¶

Derive the key of one forward pass.

Parameters:

model_fingerprint – What identifies the model, its architecture and weights.
model_args – The positional arguments of the run_with_cache call.
names_filter – The activations cached, as _normalize_names_filter returns them.
extra – Any other keyword argument that changes the result.

Returns:

A SHA-256 hex digest.

load(digest: str, model_name: str) → dict | None[source]¶

Read an entry back.

Parameters:

digest – The entry’s key.
model_name – The model the entry belongs to.

Returns:

The stored payload, on the CPU, or None when there is no such entry or it was written in another format.

store(digest: str, model_name: str, payload: dict) → None[source]¶

Write an entry atomically.

Parameters:

digest – The entry’s key.
model_name – The model the entry belongs to.
payload – What to store.

class entoscope.model.cache.CachedModel(model: Model, cache_dir: str | None = None)[source]¶

Bases: object

Transparent read-through disk cache around a transformer_lens model.

Overrides only run_with_cache, and forwards every other attribute to the wrapped model, a method that returns the model itself returning the wrapper instead. Cache keys fold in a sampled digest of the model’s weights, taken on every call, so neither a fine-tune nor an edit made after wrapping, such as enable_compatibility_mode() or a change of dtype, reads the entries of the weights it replaced.

Parameters:

model – A TransformerBridge or a HookedTransformer.
cache_dir – The directory activations are stored under. None makes the wrapper a pure pass-through.

run_with_cache(*model_args: Any, **kwargs: Any) → tuple[Any, ActivationCache][source]¶

Run the model and cache its activations, reading an earlier identical run back from disk.

A callable names_filter or incl_bwd=True cannot be keyed, so such a call always runs the model.

Parameters:

*model_args – The wrapped model’s run_with_cache positional arguments, such as tokens.
**kwargs – Its keyword arguments, such as names_filter.

Returns:

The logits and the activation cache, on the model’s device.

entoscope.model.token.escape_token(token: str) → str[source]¶

Escape a token’s non-printable characters, since raw ones such as \r break CSV round-trips.

Parameters:: token – A token string.
Returns:: The token with every non-printable character written as its escape sequence.

entoscope.model.token.str_token(model: TransformerBridge | CachedModel, token: int) → str[source]¶

Spell one token, safe to write into a CSV or XML file.

Safe to write, not to read back blindly: pandas’ default read_csv takes real tokens such as “null”, “None” or “NA” for missing values. Read string columns with keep_default_na=False to keep them verbatim, or go by the token id every collector writes beside its string.

Parameters:

model – The model whose tokenizer spells the token.
token – A token id.

Returns:

The token’s string, escaped as escape_token does.

entoscope.model.token.str_token_labels(model: TransformerBridge | CachedModel, tokens: Int[Tensor, '... pos']) → list[str][source]¶

Spell the tokens of one prompt, safe to write into a CSV or XML file.

Parameters:

model – The model whose tokenizer spells the tokens.
tokens – The prompt’s tokens, shaped [pos] or [1, pos].

Returns:

One escaped string per position.

Raises:

ValueError – If tokens holds more than one prompt.

class entoscope.model.weight.ModelWeight(*values)[source]¶

Bases: StrEnum

A weight of the model, by its HookedTransformer name.

entoscope.model.weight.activation_weight(model: TransformerBridge | CachedModel, activation: ModelLayerLocation, layer: int) → Float[Tensor, '... d_model'][source]¶

Read the weight ACTIVATION_WEIGHT names for activation.

That is the weight computing the activation off the residual stream, or, for one other activations compute (attn.z, mlp.post), the weight writing it back. A gated MLP computes mlp.pre with its gate and mlp.pre_linear with W_in, and any other MLP computes mlp.pre with W_in and has no mlp.pre_linear.

Parameters:

model – The model to read.
activation – An activation ACTIVATION_WEIGHT maps.
layer – The block to read the weight of.

Returns:

The weight, as model_weight returns it.

Raises:

ValueError – If activation is mlp.pre_linear on a model without a gated MLP.

entoscope.model.weight.model_weight(model: TransformerBridge | CachedModel, name: ModelWeight, layer: int | None = None) → Float[Tensor, '... d_model'][source]¶

Read a weight as one d_model vector per unit it reads or writes.

The layout is the same whichever side of the residual stream the weight is on, and whichever backend holds it, so a lens can read every weight the same way.

Parameters:

model – The model to read.
name – The weight.
layer – The block to read, for every weight but the embedding and unembedding.

Returns:

[d_vocab, d_model] for the embedding and unembedding, [d_mlp, d_model] for an MLP weight and [head, d_head, d_model] for an attention one.

Raises:

ValueError – If a block’s weight is read without a layer, or W_gate on a model without a gated MLP.

entoscope.model.unembed.build_unembed(model: TransformerBridge | CachedModel) → Unembed[source]¶

Build the final norm and unembedding a tuned lens reads a residual stream through.

Parameters:: model – The model, wrapped in a CachedModel or not.
Returns:: A frozen tuned-lens Unembed holding the model’s final norm and unembedding.