Steps — Record

Record step: captures activations via nnsight trace.

ActivationStore

Stores per-layer/module activations for contrastive dataset splits.

Attributes:

• positive: {Node: tensor} for positive texts, keyed by component
• address. Accepts shorthand on lookup (store.positive[5]).
• negative: {Node: tensor} for negative texts.
• position: Token-position selection applied at record time
• ("last" / "first" / "mean" / int, or "none" to keep
• every position). Sets the tensor rank: reduced modes give
• [N, d_model]; "none" gives [N, seq, d_model].
• per_head: Whether activations are split per attention head, which adds a
• trailing [..., n_heads, head_dim] pair of dims.
• positive_token_mask: [N, seq] valid-token mask for positive when
• position="none"; None otherwise.
• negative_token_mask: [N, seq] valid-token mask for negative when
• position="none"; None otherwise.

__init__

def __init__(self, positive: dict[Node, Tensor], negative: dict[Node, Tensor], position: str | int = 'last', per_head: bool = False, positive_token_mask: Tensor | None = None, negative_token_mask: Tensor | None = None) -> None:

LabeledActivationStore

Stores per-layer/module activations with associated per-example labels.

Attributes:

• activations: {Node: tensor} token-position activations, keyed by
• component address. Accepts shorthand on lookup.
• labels: tensor [N] integer labels.
• position: Token-position selection applied at record time (see
• “: class:ActivationStore).
• per_head: Whether activations are split per attention head.
• token_mask: [N, seq] valid-token mask when position="none";
• None otherwise.

__init__

def __init__(self, activations: dict[Node, Tensor], labels: Tensor, position: str | int = 'last', per_head: bool = False, token_mask: Tensor | None = None) -> None:

Record

Capture residual-stream activations via nnsight.

Reads from results:

• results['dataset']: MuranoDataset or LabeledDataset

Writes to results:

• results['record']: ActivationStore or LabeledActivationStore

Args:

• model: Wrapped model to record from.
• layers: Layer indices to record, or "all" for every layer.
• position: Token position to record at. One of "last", "first",
• "mean", an integer token index, or "none" to keep every
• position (full-position recording).
• batch_size: Forward-pass batch size; must be >= 1.
• per_head: If True, split attention activations per head, producing a
• trailing [..., n_heads, head_dim] pair of dims. Only valid for
• attention modules; the per-head signal is the input to the
• attention output projection (the concatenated head outputs).

Raises:

• ValueError: If position or batch_size is invalid, or
• layers is a string other than "all".
• NotImplementedError: If per_head is set for a module whose
• architecture’s attention output projection is not recognized.

Note:

• Full-position (position="none") and per-head recording accumulate
• larger tensors in CPU memory than reduced recording; combining them with
• layers="all", multiple modules, and large batches is memory-bound.

__init__

def __init__(self, model: ModelBackend, layers: list[int] | str = 'all', modules: str | list[str] = 'residual', position: str | int = 'last', batch_size: int = 8, per_head: bool = False):

expected_read_types

def expected_read_types(self, results = None, available_types = None):

Return {"dataset": (MuranoDataset, LabeledDataset)}.

expected_write_types

def expected_write_types(self, results = None, available_types = None):

Return the write type for record, narrowed by the upstream dataset type.

The output store type mirrors the input dataset type: MuranoDataset produces ActivationStore; LabeledDataset produces LabeledActivationStore. Falls back to the union when the dataset type is not yet known.