`torchlens.dataset_extraction`¶

Batched dataset extraction with a self-describing, resumable disk artifact.

This module is the public face of the extraction v2 runtime (extract memo items 4-19; the engine lives in :mod:torchlens._extraction). In disk mode (path=) the artifact directory is the v2 layout: a BOUNDED manifest.json (written at creation, once when batch zero freezes the plan, and at terminal status — never per shard), an APPEND-ONLY fsynced ledger.jsonl with one line per committed shard, a write-once ordered stimulus_ids.json sidecar, and immutable batch-major shards (safetensors by default; .pt stays readable forever and writable via shard_format="pt"). The commit protocol is validate -> temp shard -> flush/fsync -> atomic rename -> append/fsync ledger row.

Every batch travels the ONE manifested store pipeline::

collate -> position/mask validation -> model call -> selected output
        -> pool (captured dtype, raw axes, mask) -> postprocess
        -> row + logical-shape validation -> per-key save-dtype cast
        -> CPU-contiguous snapshot -> writer

The run signature carries MODEL IDENTITY, the collate policy, callable identity, pool/dtype/ragged policies, and the selector request — all compared FIELD BY FIELD on resume.

Read artifacts back through :func:open_extraction (the documented lazy path) or :func:load_extraction (the eager compatibility wrapper behind an exact byte guard); export with :func:export_extraction.

Every spelling introduced here is DOCUMENTED-UNSTABLE pending the naming/UI sprint.

BatchEnvelope `dataclass` ¶

One collated batch, typed.

Attributes:

Name	Type	Description
`args`	`tuple[Any, ...]`	Positional arguments for `model.forward`.
`kwargs`	`dict[str, Any]`	Keyword arguments for `model.forward`.
`num_rows`	`int`	INPUT-derived number of stimulus rows in this batch — the row-truth every selected output is validated against.
`mask`	`Tensor \| None`	The batch's two-dimensional attention mask when the collator supplied one (nonzero = token present); `None` otherwise.
`disclosure`	`dict[str, Any]`	JSON-portable collate disclosure recorded in the manifest and the resume signature (collate kind, tokenizer identity, pad-token source, ...).

Examples:

>>> envelope = BatchEnvelope(args=(x,), kwargs={}, num_rows=1, mask=None, disclosure={})
>>> envelope.num_rows
1
>>> envelope.first_input() is x
True

first_input ¶

first_input() -> Any

Return the batch's primary input container for legacy paths.

Returns:

Type	Description
`Any`	The single positional arg when there is exactly one and no kwargs, else the kwargs mapping when purely keyword-shaped, else the args tuple.

RaggedBatch `dataclass` ¶

One shard's trimmed carrier for one output key (the Arrow triplet).

Attributes:

Name	Type	Description
`values`	`Tensor`	Packed per-row slices concatenated along axis 0: `[sum(extents), *feature_shape]`.
`offsets`	`Tensor`	`int64` tensor of length `rows + 1`; row `i` occupies `values[offsets[i]:offsets[i+1]]`. The STORED authority.
`row_shapes`	`tuple[tuple[int, ...], ...]`	Per-row logical shapes (each row's true `[extent, *feature]`).

Examples:

>>> batch = RaggedBatch(
...     values=torch.zeros(5, 4), offsets=torch.tensor([0, 2, 5]), row_shapes=((2, 4), (3, 4))
... )
>>> batch.num_rows
2
>>> tuple(batch.row(1).shape)
(3, 4)

num_rows `property` ¶

num_rows: int

Number of logical rows in the carrier.

Returns:

Type	Description
`int`	`len(offsets) - 1`.

row ¶

row(index: int) -> torch.Tensor

Return one exact ragged row.

Parameters:

Name	Type	Description	Default
`index`	`int`	Row index within this carrier.	required

Returns:

Type	Description
`Tensor`	The row at its true extent.

TokenRows `dataclass` ¶

Stored keys flattened to one row per token (DOCUMENTED-UNSTABLE).

Attributes:

Name	Type	Description
`values`	`dict[str, Tensor]`	`key -> [n_tokens, *feature_shape]` per selected key, every key aligned token for token.
`stimulus_row`	`Tensor`	`int64 [n_tokens]`: the global stimulus row of each token.
`position`	`Tensor`	`int64 [n_tokens]`: each token's index within its stimulus's true (unpadded) extent.

Examples:

>>> import tempfile
>>> def collate(texts):
...     ids = torch.tensor([[int(word) for word in text.split()] for text in texts])
...     return {"input_ids": ids, "attention_mask": torch.ones_like(ids)}
>>> class Embed(nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.emb = nn.Embedding(10, 4)
...     def forward(self, input_ids, attention_mask=None):
...         return self.emb(input_ids)
>>> with tempfile.TemporaryDirectory() as folder:
...     _ = tl.dataset_extraction.extract_dataset(
...         Embed().eval(), ["1 2", "3 4 5"], {"h": "emb"}, batch_size=1, path=folder,
...         collate=collate, ragged="trim", progress=False,
...     )
...     tokens = tl.dataset_extraction.open_extraction(folder).token_rows()
>>> tokens.num_tokens, tokens.position.tolist()
(5, [0, 1, 0, 1, 2])

num_tokens `property` ¶

num_tokens: int

Return the number of token rows.

Returns:

Type	Description
`int`	`len(stimulus_row)`.

LoadedExtraction `dataclass` ¶

A dataset-extraction artifact read back with its self-description.

Attributes:

Name	Type	Description
`manifest`	`dict[str, Any]`	Parsed `manifest.json` document (site identity, stimulus provenance, axis semantics, dtypes, devices, run signature, TorchLens version).
`activations`	`dict[str, Any]`	Materialized payloads keyed by output key: concatenated tensors for dense keys, one merged :class:`~torchlens.dataset_extraction.RaggedBatch` for trimmed keys (never silent padding, never object arrays).
`batch_paths`	`list[Path]`	Shard files in ledger consumption order.

Examples:

>>> import tempfile
>>> with tempfile.TemporaryDirectory() as tmp:
...     stimuli = torch.randn(4, 3, 8, 8)
...     _ = tl.extract_dataset(model, stimuli, ["1"], path=tmp, progress=False)
...     loaded = tl.dataset_extraction.load_extraction(tmp)
...     shape = tuple(loaded.activations["1"].shape)
>>> isinstance(loaded, LoadedExtraction), shape
(True, (4, 4, 8, 8))

input_preprocessing `property` ¶

input_preprocessing: dict[str, Any]

The manifest's input-preprocessing block (legacy reads as unknown).

Returns:

Type	Description
`dict[str, Any]`	See :func:`input_preprocessing_of`.

export_extraction ¶

export_extraction(path: str | Path, dest: str | Path, *, file_format: str = 'npy', keys: list[str] | None = None, widen_unsupported: bool = True) -> Path

Export one COMPLETE extraction artifact (read-only, streaming).

Parameters:

Name	Type	Description	Default
`path`	`str \| Path`	Source artifact directory.	required
`dest`	`str \| Path`	Destination directory (created; published atomically by rename).	required
`file_format`	`str`	`"npy"` (base install) \| `"hdf5"` \| `"mat"` (both behind the named optional extra; lazy imports).	`'npy'`
`keys`	`list[str] \| None`	Optional output-key subset.	`None`
`widen_unsupported`	`bool`	Widen bf16/fp8 members to fp32 with disclosure (default) or refuse typed.	`True`

Returns:

Type	Description
`Path`	The published destination directory.

Raises:

Type	Description
`InvalidArgumentError`	`extraction_export_format_invalid` outside the format vocabulary; `extraction_export_dependency_missing` when the extra is absent; `extraction_export_space_insufficient` on the preflight; `extraction_export_dtype_unsupported` on refused widening; `extraction_export_mat_v5_limit` for MAT members over 2 GB.

Examples:

>>> import tempfile
>>> from pathlib import Path
>>> stimuli = [torch.randn(3, 8, 8) for _ in range(2)]
>>> with tempfile.TemporaryDirectory() as src, tempfile.TemporaryDirectory() as out:
...     _ = tl.dataset_extraction.extract_dataset(
...         model, stimuli, {"relu": "1"}, batch_size=2, path=src, progress=False
...     )
...     dest = export_extraction(src, Path(out) / "npy", file_format="npy")
...     print(dest.name, dest.is_dir())
npy True

hf_collate ¶

hf_collate(tokenizer: Any = None, *, pad_token: str | None = None, **tokenizer_kwargs: Any) -> Callable[[list[Any], nn.Module, dict[str, Any]], BatchEnvelope]

Build the HF text collate helper (DOCUMENTED-UNSTABLE).

Parameters:

Name	Type	Description	Default
`tokenizer`	`Any`	Caller-supplied tokenizer (first in the resolution order); `None` defers to the attached-then-hub doors.	`None`
`pad_token`	`str \| None`	`None` (a missing pad token refuses typed) or `"eos"` — the explicit, disclosed opt-in recorded as `pad_token_source`.	`None`
`**tokenizer_kwargs`	`Any`	Extra tokenizer call kwargs (`truncation=`, `max_length=`, ...).	`{}`

Returns:

Type	Description
`Callable`	An engine collate callable: `(items, model, run_state) -> BatchEnvelope`.

Examples:

>>> collate = hf_collate(pad_token="eos", truncation=True, max_length=16)
>>> callable(collate)
True

open_extraction ¶

open_extraction(path: str | Path, *, verify: str = 'structural', in_progress: bool = False) -> ExtractionReader

Open one extraction artifact lazily (the documented read path).

Opening reads bounded manifest/ledger metadata only; payload bytes move on demand through the reader.

Parameters:

Name	Type	Description	Default
`path`	`str \| Path`	Artifact directory written by disk-mode `extract_dataset`.	required
`verify`	`str`	`"structural"` (default: manifest/ledger well-formed, files present at their exact ledgered byte sizes) or `"first_access"` (each shard's CRC verified once per handle lifetime).	`'structural'`
`in_progress`	`bool`	Open a still-writing artifact's trusted committed prefix; `reader.refresh()` rescans.	`False`

Returns:

Type	Description
`ExtractionReader`	The lazy reader.

Raises:

Type	Description
`DatasetExtractionResumeError`	On a missing/invalid manifest or an incomplete artifact opened without `in_progress=True`.
`InvalidArgumentError`	`extraction_reader_verify_invalid` outside the closed verification vocabulary.

Examples:

>>> import tempfile
>>> stimuli = [torch.randn(3, 8, 8) for _ in range(4)]
>>> with tempfile.TemporaryDirectory() as folder:
...     _ = tl.dataset_extraction.extract_dataset(
...         model, stimuli, {"relu": "1"}, batch_size=2, path=folder, progress=False
...     )
...     reader = open_extraction(folder)
...     print(reader.keys, reader.n_stimuli, tuple(reader.row(0)["relu"].shape))
['relu'] 4 (4, 8, 8)

register_pure_module ¶

register_pure_module(name: str) -> None

Register a module namespace as pure for callable classification.

Registration is an AUDITED ASSERTION, not an invisible trust grant: the module name and its installed version are recorded in every classification record that folds it, and those records ride the extraction manifest.

Parameters:

Name	Type	Description	Default
`name`	`str`	Top-level module name (e.g. `"einops"`).	required

Raises:

Type	Description
`InvalidArgumentError`	`extraction_pure_module_invalid` when the name is empty, dotted, or not an importable installed module.

Examples:

>>> try:
...     register_pure_module("numpy.linalg")
... except tl.errors.InvalidArgumentError as err:
...     print(type(err).__name__)
InvalidArgumentError

extract_dataset ¶

extract_dataset(model: Module, stimuli: Any, layers: Any, batch_size: int = 32, device: device | str | None = None, path: str | Path | None = None, output_transform: Any = None, progress: bool = True, *, resume: bool = False, stimulus_ids: Iterable[str] | None = None, model_identity: Any = 'measured', input_transform: Any = None, input_provenance: Any = None, collate: Any = None, pool: Any = None, dtype: Any = None, ragged: str = 'refuse', shard_format: str | None = None, checksums: str = 'fast', pipeline_id: str | None = None, intervene: Any = None, store_input_ids: bool = False) -> dict[str, torch.Tensor] | list[Path]

Extract outputs from an iterable dataset in batches.

Row i of every stored tensor corresponds to stimulus i in iteration order; shards are consumed in ledger order.

Parameters:

Name	Type	Description	Default
`model`	`Module`	PyTorch model to run.	required
`stimuli`	`Any`	Tensor with a leading batch dimension, or an iterable of stimulus items (tensors, containers, or TEXT strings — text routes through the once-per-run tokenizer with padding and the full mapping forwarded).	required
`layers`	`Any`	List or mapping of layer lookups (:func:`torchlens.extract` semantics), a capture SELECTOR (`tl.func`/`tl.in_module`/ composed predicates), or a selector-valued mapping. Selector requests resolve against batch zero in execution order, FREEZE (ordered site keys + pass-qualified labels), and every later batch must attest the exact plan.	required
`batch_size`	`int`	Number of stimuli per forward pass.	`32`
`device`	`device \| str \| None`	Optional device for model and stimuli.	`None`
`path`	`str \| Path \| None`	Optional directory. When supplied, each batch is committed as an immutable shard under the v2 protocol and paths are returned.	`None`
`output_transform`	`Any`	Transform slot, through the ONE :func:`torchlens.transforms. coerce_transform` door (DOCUMENTED-UNSTABLE). Selector-valued `layers=` accepts `None` or one global chain (per-site mappings need knowable keys).	`None`
`progress`	`bool`	Whether to wrap batch iteration with `tqdm`.	`True`
`resume`	`bool`	Disk mode only (DOCUMENTED-UNSTABLE): continue an interrupted run from its trusted ledger prefix. Every semantic signature field is compared individually; the skipped prefix is replayed at zero forward cost and verified against the ledgered per-shard `input_digest` where recorded.	`False`
`stimulus_ids`	`Iterable[str] \| None`	Optional per-stimulus identifiers (DOCUMENTED-UNSTABLE): non-empty strings, duplicates legal (recorded as `ids_unique`); validated against cardinality BEFORE any forward for sized inputs, in lockstep for unsized iterables, and with exact equality before completion. Disk mode only.	`None`
`model_identity`	`Any`	Disk mode (DOCUMENTED-UNSTABLE): `"measured"` (default — cryptographic threaded-Merkle state digest), `"none"` (recorded opt-out), or an assertion Mapping for unmeasurable models.	`'measured'`
`input_transform`	`Any`	The unambiguous INPUT-preprocessing path (tvscope B5, DOCUMENTED-UNSTABLE spelling), distinct from `output_transform=` which transforms OUTPUTS before storage. Pass a :class:`torchlens.preprocessing.Resolution` (the golden path: the authority's own callable and its provenance travel together, the manifest block is stamped verified-by-construction, and the transform identity joins the resume signature) or a bare callable (applied and disclosed as opaque; the artifact becomes non-resumable across interruptions, mirroring opaque output transforms). Item-level callables run per stimulus and stack; batch-native processors run once per batch.	`None`
`input_provenance`	`Any`	Provenance-only stamp (tvscope B4, DOCUMENTED-UNSTABLE spelling) for stimuli preprocessed OUTSIDE this run: a :class:`torchlens.preprocessing.PreprocessingComparison` (its verdict rides the manifest block), a `Resolution`, or a `ResolvedPreprocessing` record (authority-only; verdict stays `unknown` with `applied_not_audited` disclosed). Disk mode only: the in-memory result carries no manifest, so a provenance stamp there is a false affordance and refuses typed.	`None`
`collate`	`Any`	Optional collate (DOCUMENTED-UNSTABLE): a callable receiving the batch's item list and returning a :class:`~torchlens.dataset_extraction.BatchEnvelope`, a Mapping (kwargs), or a Tensor (one positional arg) — bare tuples refuse as ambiguous. :func:`torchlens.dataset_extraction.hf_collate` builds the HF text helper (tokenizer-once, pad-token refusals, full mapping).	`None`
`pool`	`Any`	Optional pool preset (DOCUMENTED-UNSTABLE): one of `flatten` / `spatial_mean` / `spatial_max` / `token_mean` / `token_max` / `token_sum` / `cls` / `last_token` (token presets are MASK-AWARE), a `{"preset": ..., "unmasked": true}` recorded override, or a per-key mapping. Runs in captured dtype on raw axes BEFORE transforms; the documented raggedness-killer.	`None`
`dtype`	`Any`	Optional save dtype (DOCUMENTED-UNSTABLE): a `torch.dtype`/name or per-key mapping; `None` preserves. Cast after pool/postprocess, before the snapshot (the host copy in disk mode); bf16/fp8 store byte-exact in safetensors shards.	`None`
`ragged`	`str`	`"refuse"` (default) \| `"trim"` \| `"as_captured"`. `refuse` fails typed BEFORE committing the first width-drifted shard; `trim` stores true per-stimulus extents as the packed values/offsets/shapes carrier (mask-derived `(start, extent)` slicing — left padding is why extent alone mis-slices); `as_captured` is reserved (refuses typed until its point release).	`'refuse'`
`shard_format`	`str \| None`	`None` (default: safetensors on fresh runs, the recorded format on resume) \| `"safetensors"` \| `"pt"`.	`None`
`checksums`	`str`	The integrity level: `"fast"` (final-file CRC-32 plus the on-device value reduction; default), `"crypto"` (blake2b file digests), `"none"` (recorded opt-out; verification refuses).	`'fast'`
`pipeline_id`	`str \| None`	Caller-asserted pipeline identity: permits resuming past a PARTIAL callable classification or a deliberate complete-class refactor, with strict continuity — the id must already be recorded in the artifact's prefix.	`None`
`intervene`	`Any`	Optional :class:`torchlens.intervention._spec.InterventionSpec` (`tl.when(selector, helper)`; DOCUMENTED-UNSTABLE): every batch is captured with the spec armed, so stored rows are the INTERVENED activations. Batch zero freezes the structural sites the spec fired at; a batch firing nowhere or elsewhere refuses typed. Disk mode records the counterfactual in the manifest's `tl_counterfactual_harvest_v1` block (address class, sites, fire receipt digest) and ledgers each shard's fire receipt; the intervention identity joins the resume signature. Module- and op-type-addressed specs run; finer address classes refuse typed. Helpers must declare `is_batch_independent=True`. String or mapping `layers=` only.	`None`
`store_input_ids`	`bool`	Store each batch's token ids as the ordinary `input_ids` key (DOCUMENTED-UNSTABLE, token store), trimmed by the batch mask under `ragged="trim"` like the token-axis activations; read token rows back with `open_extraction(...).token_rows()`.	`False`

Returns:

Type	Description
`dict[str, Tensor] \| list[Path]`	In-memory outputs on the model's device (no host copy), or shard paths.

Raises:

Type	Description
`InvalidArgumentError`	On closed-vocabulary violations, in-memory false affordances, left-padding that cannot be corrected, row-count drift, ragged refusals, pool/dtype geometry refusals, and id validation.
`DatasetExtractionResumeError`	When the artifact in `path` cannot be safely continued.
`TransformContractError`	When the transform slot cannot be coerced or violates its plan.

Notes

Every forward runs under torch.no_grad() with the model in eval mode, and every submodule's exact training flag is restored afterward (exception paths included).

Examples:

>>> stimuli = torch.randn(4, 3, 8, 8)
>>> out = tl.extract_dataset(model, stimuli, ["1"], batch_size=2, progress=False)
>>> tuple(out["1"].shape)
(4, 4, 8, 8)
>>> import tempfile
>>> with tempfile.TemporaryDirectory() as tmp:
...     shards = tl.extract_dataset(
...         model, stimuli, ["1"], batch_size=2, path=tmp, progress=False
...     )
...     len(shards)
2

load_extraction ¶

load_extraction(path: str | Path, layers: Iterable[str] | None = None, *, max_bytes: int | None = None) -> LoadedExtraction

Eagerly load a disk-mode extraction artifact (the guarded wrapper).

This is a thin materialization over :func:open_extraction — the documented lazy path — behind an EXACT byte guard computed from ledger facts before any allocation.

Parameters:

Name	Type	Description	Default
`path`	`str \| Path`	Directory previously written by disk-mode :func:`extract_dataset`.	required
`layers`	`Iterable[str] \| None`	Optional subset of output keys to load; defaults to every key.	`None`
`max_bytes`	`int \| None`	Optional explicit byte-budget override for the guard (`None` applies the default: half of measurable available host memory).	`None`

Returns:

Type	Description
`LoadedExtraction`	Manifest, materialized activations, and shard paths.

Raises:

Type	Description
`DatasetExtractionResumeError`	If the manifest is missing/invalid or the artifact is incomplete.
`InvalidArgumentError`	`extraction_eager_budget_exceeded` when the exact requested bytes exceed the budget (the refusal names the reader and the override); `extraction_reader_key_unknown` for unknown keys.

relabel_extraction ¶

relabel_extraction(path: str | Path, new_ids: Iterable[str]) -> dict[str, Any]

Relabel a COMPLETE artifact's stimulus ids (the audited verb).

Relabeling is deliberately OUT of resume — the sidecar is write-once there — and produces a NEW artifact identity: the manifest's id digest changes, the prior digest is recorded, and an audit line rides the manifest permanently. Per-shard ledger ids_range_digest facts keep describing the WRITE-TIME labeling (the ledger is append-only); the audit trail is the pointer between the two.

Parameters:

Name	Type	Description	Default
`path`	`str \| Path`	Artifact directory.	required
`new_ids`	`Iterable[str]`	Replacement identifiers, one per stimulus row, in row order.	required

Returns:

Type	Description
`dict[str, Any]`	The appended audit record.

Raises:

Type	Description
`InvalidArgumentError`	`extraction_relabel_invalid` when the artifact is not complete or the id list mislengths; `extraction_stimulus_ids_invalid` for malformed identifiers.

Examples:

>>> import tempfile
>>> with tempfile.TemporaryDirectory() as tmp:
...     stimuli = torch.randn(2, 3, 8, 8)
...     _ = tl.extract_dataset(
...         model, stimuli, ["1"], path=tmp, progress=False, stimulus_ids=["a", "b"]
...     )
...     audit = relabel_extraction(tmp, ["first", "second"])
...     ids = tl.dataset_extraction.open_extraction(tmp).stimulus_ids()
>>> audit["kind"], ids
('relabel', ['first', 'second'])

torchlens.dataset_extraction¶

BatchEnvelope dataclass ¶

first_input ¶

RaggedBatch dataclass ¶

num_rows property ¶

row ¶

TokenRows dataclass ¶

num_tokens property ¶

LoadedExtraction dataclass ¶

input_preprocessing property ¶

export_extraction ¶

hf_collate ¶

open_extraction ¶

register_pure_module ¶

extract_dataset ¶

load_extraction ¶

relabel_extraction ¶

`torchlens.dataset_extraction`¶

BatchEnvelope `dataclass` ¶

RaggedBatch `dataclass` ¶

num_rows `property` ¶

TokenRows `dataclass` ¶

num_tokens `property` ¶

LoadedExtraction `dataclass` ¶

input_preprocessing `property` ¶